Require Export header_extensible.

From MetaCoq.Template Require Export All.
From MetaCoq.Checker Require Export Checker uGraph.

Local Set Implicit Arguments.
Require Export String List Omega.
Export ListNotations.
Local Unset Strict Implicit.

From MetaCoq.Template Require Export monad_utils.
Export MonadNotation.

Class subtermC A := subterm_rel : A -> A -> Prop.
Class InRelC B := {in_subtype : Type ; in_rel : in_subtype -> B -> Prop }.

Fixpoint destArity Γ (t : term) :=
  match t with
  | tProd na t b => destArity (Γ ,, vass na t) b
  | tLetIn na b b_ty b' => destArity (Γ ,, vdef na b b_ty) b'
  | s => (Γ, s)
  end.

Derive Subterm for nat.

Fixpoint replace_term (c : term) d (t : term) :=
  if @eq_term config.type_in_type init_graph c t then d else
    match t with
  | tRel i => tRel i
  | tEvar ev args => tEvar ev (List.map (replace_term c d) args)
  | tLambda na T M => tLambda na (replace_term c d T) (replace_term c d M)
  | tApp u v => tApp (replace_term c d u) (List.map (replace_term c d) v)
  | tProd na A B => tProd na (replace_term c d A) (replace_term c d B)
  | tCast C kind t => tCast (replace_term c d C) kind (replace_term c d t)
  | tLetIn na b t b' => tLetIn na (replace_term c d b) (replace_term c d t) (replace_term c d b')
  | tCase ind p C brs =>
    let brs' := List.map (on_snd (replace_term c d)) brs in
    tCase ind (replace_term c d p) (replace_term c d C) brs'
  | tProj p C => tProj p (replace_term c d C)
  | tFix mfix idx =>
    let mfix' := List.map (map_def (replace_term c d) (replace_term c d)) mfix in
    tFix mfix' idx
  | tCoFix mfix idx =>
    let mfix' := List.map (map_def (replace_term c d) (replace_term c d)) mfix in
    tCoFix mfix' idx
  | tConst name u => tConst name u
  | x => x
  end.
Require Import Ascii.

Fixpoint name_after_dot' (s : string) (r : string) :=
  match s with
  | EmptyString => r
  | String "#" xs => name_after_dot' xs xs (* see Coq_name in a section *)
  | String ("."%char) xs => name_after_dot' xs xs
  | String _ xs => name_after_dot' xs r
  end.

Definition name_after_dot s := name_after_dot' s s.

Fixpoint fixNames (t : term) :=
  match t with
  | tRel i => tRel i
  | tEvar ev args => tEvar ev (List.map (fixNames) args)
  | tLambda na T M => tLambda na (fixNames T) (fixNames M)
  | tApp u v => tApp (fixNames u) (List.map (fixNames) v)
  | tProd na A B => tProd na (fixNames A) (fixNames B)
  | tCast C kind t => tCast (fixNames C) kind (fixNames t)
  | tLetIn na b t b' => tLetIn na (fixNames b) (fixNames t) (fixNames b')
  | tCase ind p C brs =>
    let brs' := List.map (on_snd (fixNames)) brs in
    tCase ind (fixNames p) (fixNames C) brs'
  | tProj p C => tProj p (fixNames C)
  | tFix mfix idx =>
    let mfix' := List.map (map_def (fixNames) (fixNames)) mfix in
    tFix mfix' idx
  | tCoFix mfix idx =>
    let mfix' := List.map (map_def (fixNames) (fixNames)) mfix in
    tCoFix mfix' idx
  | tConst name u => tConst (name_after_dot name) u
  | tInd (mkInd name i) u => tInd (mkInd (name_after_dot name) i) u
  | x => x
  end.

Fixpoint replace_const (c : kername) d (t : term) :=
  match t with
  | tRel i => tRel i
  | tEvar ev args => tEvar ev (List.map (replace_const c d) args)
  | tLambda na T M => tLambda na (replace_const c d T) (replace_const c d M)
  | tApp u v => tApp (replace_const c d u) (List.map (replace_const c d) v)
  | tProd na A B => tProd na (replace_const c d A) (replace_const c d B)
  | tCast C kind t => tCast (replace_const c d C) kind (replace_const c d t)
  | tLetIn na b t b' => tLetIn na (replace_const c d b) (replace_const c d t) (replace_const c d b')
  | tCase ind p C brs =>
    let brs' := List.map (on_snd (replace_const c d)) brs in
    tCase ind (replace_const c d p) (replace_const c d C) brs'
  | tProj p C => tProj p (replace_const c d C)
  | tFix mfix idx =>
    let mfix' := List.map (map_def (replace_const c d) (replace_const c d)) mfix in
    tFix mfix' idx
  | tCoFix mfix idx =>
    let mfix' := List.map (map_def (replace_const c d) (replace_const c d)) mfix in
    tCoFix mfix' idx
  | tConst name u => if eq_string name c then d else tConst name u
  | x => x
  end.

Fixpoint remove_injs (k : nat) (u : term) {struct u} : term :=
  match u with
  | tRel n => tRel n
  | tEvar ev args => tEvar ev (map (remove_injs k) args)
  | tCast c kind ty => tCast (remove_injs k c) kind (remove_injs k ty)
  | tProd na A B => tProd na (remove_injs k A) (remove_injs (S k) B)
  | tLambda na T M => tLambda na (remove_injs k T) (remove_injs (S k) M)
  | tLetIn na b ty b' => tLetIn na (remove_injs k b) (remove_injs k ty) (remove_injs (S k) b')
  | tApp u0 v => let L := remove_injs k u0 in
                let R := map (remove_injs k) v in
                match L, R with tConst "header_extensible.retract_I" _ , [_;_;_;tRel k'] => (* if k =? k' then *) tRel k' (* else mkApps L R *)
                           | _,_ => mkApps L R end
  | tCase ind p c brs => let brs' := map (on_snd (remove_injs k)) brs in tCase ind (remove_injs k p) (remove_injs k c) brs'
  | tProj p c => tProj p (remove_injs k c)
  | tFix mfix idx => let k' := #|mfix| + k in let mfix' := map (map_def (remove_injs k) (remove_injs k')) mfix in tFix mfix' idx
  | tCoFix mfix idx =>
      let k' := #|mfix| + k in let mfix' := map (map_def (remove_injs k) (remove_injs k')) mfix in tCoFix mfix' idx
  | _ => u
  end.

Fixpoint replace_terms Repl t :=
  match Repl with
  | ((c, s) :: Repl) => replace_terms Repl (replace_term c s t)
  | _ => t
  end.

Fixpoint replace_consts Repl t :=
  match Repl with
  | ((tConst na u, s) :: Repl) => replace_consts Repl (replace_const na s t)
  | _ => t
  end.

Fixpoint replace_ext T T_ext t :=
  match t with
  | tProd na argT retT => if @eq_term config.default_checker_flags init_graph argT T_ext
                         then tProd na T (remove_injs 0 (replace_ext T T_ext retT))
                         else tProd na argT (replace_ext T T_ext retT)
  | s => s
  end.

(* todo here  *)
Definition genIH retT decl_list LT (T T_ext : term) Repl :=
  match LT with
    Some LT =>
    let IH := replace_terms Repl retT in
    let IH := (it_mkProd_or_LetIn decl_list ((tProd nAnon (tApp LT [tRel 0; tRel 1]) (lift 1 0 IH)))) in
    let IH := replace_ext T T_ext IH in
    IH
  | None =>
    let IH := replace_terms Repl retT in
    let IH := (it_mkProd_or_LetIn decl_list (IH)) in
    let IH := replace_ext T T_ext IH in
    IH
  end.

Fixpoint genStatement LT decl_list (T T_ext : term) Repl n B :=
  match n, B with
  | 0, retT => let IH := genIH retT (decl_list) (Some LT) T T_ext Repl in
                            it_mkProd_or_LetIn (decl_list) ((tProd nAnon IH (lift 1 0 retT)))
  | S n, tProd na argT retT => (genStatement LT (vass na argT :: decl_list) T T_ext Repl n retT)
  | _, _ => tVar "no"
  end.

Fixpoint genStatement_no_lt decl_list (T T_ext : term) Repl n B :=
  match n, B with
  | 0, retT => let IH := genIH retT (decl_list) None T T_ext Repl in
                            (IH, it_mkProd_or_LetIn (decl_list) ((tProd nAnon IH (lift 1 0 retT))))
  | S n, tProd na argT retT => (genStatement_no_lt (vass na argT :: decl_list) T T_ext Repl n retT)
  | _, _ => (tVar "no", tVar "no")
  end.

Fixpoint remove_suffix (a1 a2 s : string) : string :=
  match s with
  | EmptyString => a1
  | String "_"%char s =>
    remove_suffix (append a1 a2) (String "_"%char EmptyString) s
  | String c s => remove_suffix a1 (append a2 (String c EmptyString)) s
  end.

Definition mkLemma name (t : term) : TemplateMonad unit :=
  tmBind (tmUnquoteTyped Prop (fixNames t)) (fun t => tmBind (tmLemma name t) (fun _ => tmReturn tt)).

Definition mkVariable name (t : term) : TemplateMonad unit :=
  A <- tmAbout name ;;
  match A with None =>
               tmBind (tmUnquoteTyped Type (fixNames t)) (fun t => tmBind (tmVariable name t) (fun _ => tmReturn tt))
          | Some s => tmPrint "variable already exists, not defining"
  end.

Definition Forall' (name : ident) (n : nat) T T_ext Repl (A_P : term) (P : term) : TemplateMonad term :=
  match destArity nil P with (Gamma, _)=>
  match nth_error (rev Gamma) (Nat.pred n) with Some T' => let T' := fixNames (decl_type T') in
         T' <- tmUnquoteTyped Type T' ;;
         O <- tmInferInstance None (InRelC T') ;;
         match O with
         | Some Cl => let lt := @in_rel _ Cl in
                     LT <- tmQuote lt ;;
                     let St := genStatement LT [] T T_ext Repl n P in ret St
         | None => O <- tmInferInstance None (subtermC T') ;;
                    match O with
                    | Some Cl => let lt := @subterm_rel _ Cl in
                                LT <- tmQuote lt ;;
                                let St := genStatement LT [] T T_ext Repl n P in ret St
                    | None => let (IH, St) := genStatement_no_lt [] T T_ext Repl n P in
                             newname <- tmEval cbv (remove_suffix EmptyString EmptyString name) ;;
                             newname <- tmFreshName newname ;;
                             mkVariable newname (fixNames IH);;
                             ret St
                    end
         end
  | _ => tmFail "not enough arguments"
  end end.

Inductive genList : Type := genNil | genCons (A B : Type) (a : A) (b : B) : genList -> genList.

Fixpoint quote_list L1 : TemplateMonad (list (term * term)) :=
  match L1 with
  | genNil => ret []
  | genCons _ _ a b L => t <- tmQuote a ;; t' <- tmQuote b ;; L <- quote_list L ;; ret ((t,t') :: L)
  end.

Definition Forall (name : ident) (T : Type) (T_ext : Type) (n : nat) Repl {A_P : Type} (P : A_P) :=
  T <- tmQuote T ;;
  T_ext <- tmQuote T_ext ;;
  A_P <- tmQuote A_P ;;
  P <- tmQuote P ;;
  Repl <- quote_list Repl ;;
  St <- Forall' name n T T_ext Repl A_P P ;;
  (* (tmEval cbn St >>= tmPrint) ;; *)
  mkLemma name St.

Definition getName X (x : X) :=
  x <- tmEval cbv x;;
  t <- tmQuote x ;; match t with tLambda (nNamed na) _ _ => ret na | _ => ret "" end.

Definition ForallN (name : nat -> nat) (T : Type) (T_ext : Type) (n : nat) Repl {A_P : Type} (P : A_P) :=
  name <- getName name ;;
  T <- tmQuote T ;;
  T_ext <- tmQuote T_ext ;;
  A_P <- tmQuote A_P ;;
  P <- tmQuote P ;;
  Repl <- quote_list Repl ;;
  St <- Forall' name n T T_ext Repl A_P P ;;
  mkLemma name St.

Hint Extern 0 (subterm_rel _ _) => hnf.

Global Obligation Tactic := idtac.

Definition cns {X Y} '(x,y) := @genCons X Y x y.

Notation "x ~> y" := (x,y) (only parsing, at level 60).

Notation "[~ x ~]" := (cns x genNil) : list_scope.
Notation "[~ x ; y ; .. ; z ~]" := (cns x (cns y .. (cns z genNil) ..)) : list_scope.

Notation "'Modular' 'Lemma' na 'where' T_ext 'extends' T 'at' n 'with' C ':' P" := (ForallN (fun na : nat => na) T T_ext n C P) (at level 1, n at next level, C at next level, P at next level).
Notation "'Modular' 'Lemma' na 'where' T_ext 'extends' T 'with' C ':' P" := (ForallN (fun na : nat => na) T T_ext 0 C P) (at level 1, C at next level, P at next level).

Instance nat_subterm' : subtermC nat := lt.

Ltac inv H2 := inversion H2; subst; clear H2.

Definition tmMkDefinition name (t : term) : TemplateMonad unit :=
  tmBind (tmUnquote (fixNames t)) (fun t => tmBind (tmDefinitionRed name (Some hnf) t) (fun _ => tmReturn tt)).

Fixpoint split_forall_impl decls T :=
  match T with
  | tProd nAnon H1 H2 => (decls, H1, H2)
  | tProd na A B => split_forall_impl (vass na A :: decls) B
  | _ => ([], tVar "no", tVar "no")
  end.

Definition buildImp args H1 H2 :=
  it_mkProd_or_LetIn args (tApp (tConst "Imp" Instance.empty) [ H1; H2 ]).

Inductive GenList : Type := GenNil | GenCons (A : Type) (a : A) : GenList -> GenList.

Ltac apply_one L :=
  match constr:(L) with
  | GenNil => fail 1
  | GenCons ?a ?L => (now eapply a; eauto) || apply_one L
  end.

Class has_features (X : string) := features : list string.
Definition get_features X {H : has_features X} := features.

Definition tmTryInferInstance (red : option reductionStrategy) A :=
  I <- tmInferInstance red A ;; match I with Some x => ret x | _ => tmFail "no instance found" end.

Definition get_name_of (t : term) :=
  match t with
  | tConst c u => c
  | tInd c u => inductive_mind c
  | tApp (tConst c u) _ => c
  | tApp (tInd c u) _ => inductive_mind c
  | _ => ""
  end.

Fixpoint get_name (n : nat) (P : term) : string :=
  match n, P with
  | 0, tProd na A B => get_name_of A
  | S n, tProd na A B => get_name n B
  | _ , _ => "no name found"
  end.

Definition get_lemmas (X : string) (name : string) : TemplateMonad GenList :=
  I <- tmTryInferInstance None (has_features X) ;;
  feats <- tmEval hnf (@get_features X I) ;;
  @monad_fold_left _ TemplateMonad_Monad GenList string (fun L feature => A <- tmUnquote (tConst (append (append name "_") feature) Instance.empty) ;;
                                                                       A' <- tmEval cbn (my_projT2 A);;
                                 ret (GenCons A' L)) feats GenNil.

Definition get_lemmas_and_name (na : nat -> nat) (n : nat) (P : Type) :=
  na <- getName na;;
  P_syn <- tmQuote P ;;
  N <- tmEval cbv (name_after_dot (get_name n P_syn)) ;;
  get_lemmas N na.

Definition compose_fixpoint (na : nat -> nat) (P : Type) (body : P) :=
  na <- getName na;;
  tmDefinition na body.

Ltac int_dest n :=
  match constr:(n) with
  | 0 => let s := fresh "s" in intros s; destruct s
  | S ?n => let s := fresh "s" in intros s; int_dest n
  end.

Ltac int_inv n :=
  match constr:(n) with
  | 0 => let s := fresh "s" in intros s; inversion s
  | S ?n => let s := fresh "s" in intros s; int_inv n
  end.

Ltac int_ind ind n :=
  match constr:(n) with
  | 0 => let s := fresh "s" in intros s; induction s using ind; cbn
  | S ?n => let s := fresh "s" in intros s; int_ind ind n
  end.

Ltac fix_nat f n := let f := fresh "f" in
  match constr:(n) with
  | 0 => fix f 0
  | 1 => fix f 1
  | 2 => fix f 2
  | 3 => fix f 3
  | 4 => fix f 4
  | 5 => fix f 5
  | 6 => fix f 6
  | 7 => fix f 7
  | _ => fail
  end.

Notation "'Compose' 'Fixpoint' nm 'on' n ':' P" :=
  (let na := fun nm => nm in @compose_fixpoint na (P%type) ((ltac:(let k x := (fix_nat f (S n); int_dest n; apply_one x) in run_template_program (get_lemmas_and_name na n P) k)))) (at level 1, n at next level, P at next level).
Notation "'Compose' 'Lemma' nm 'on' n ':' P" :=
  (let na := fun nm => nm in @compose_fixpoint na (P%type) ((ltac:(let k x := (fix_nat f (S n); int_dest n; apply_one x) in run_template_program (get_lemmas_and_name na n P) k)))) (at level 1, n at next level, P at next level).
Notation "'Compose' 'Lemma' nm 'on' n 'by' 'inversion' ':' P" :=
  (let na := fun nm => nm in @compose_fixpoint na (P%type) ((ltac:(let k x := (fix_nat f (S n); int_inv n; apply_one x) in run_template_program (get_lemmas_and_name na n P) k)))) (at level 1, n at next level, P at next level).
Notation "'Compose' 'Lemma' nm 'on' n 'using' ind ':' P" :=
  (let na := fun nm => nm in @compose_fixpoint na (P%type) ((ltac:(let k x := (fix_nat f (S n); int_ind ind n; apply_one x) in run_template_program (get_lemmas_and_name na n P) k)))) (at level 1, n at next level, ind at next level, P at next level).

(* Modular Fixpoints *)

Definition genStatement_Fix (T T_ext : term) retT :=
  let IH := genIH retT [] None T T_ext [] in
    (IH, retT).

Definition mkDefinitionType name (T : Type) (t : term) : TemplateMonad unit :=
  tmBind (tmUnquoteTyped T (fixNames t)) (fun t => tmBind (tmDefinition name t) (fun _ => tmReturn tt)).

Definition ModularFixpointN (name : nat -> nat) (T : Type) (T_ext : Type) (P: Type) (A : Type) (body : A -> P) :=
  name <- getName name ;;
  T <- tmQuote T ;;
  T_ext <- tmQuote T_ext ;;
  P_syn <- tmQuote P ;;
  body_syn <- tmQuote body ;;
  let IH := genIH P_syn [] None T T_ext [] in
  newname <- tmEval cbv (remove_suffix EmptyString EmptyString name) ;;
  newname <- tmFreshName newname ;;
  mkVariable newname (fixNames IH);;
  mkDefinitionType name P (tApp body_syn [tVar newname]).

Definition ModularFixpoint2 (name : nat -> nat) (T : Type) (T_ext : Type) (P: Type) (body : P) :=
  name <- getName name ;;
  T <- tmQuote T ;;
  T_ext <- tmQuote T_ext ;;
  P_syn <- tmQuote P ;;
  body_syn <- tmQuote body ;;
  let IH := genIH P_syn [] None T T_ext [] in
  newname <- tmEval cbv (remove_suffix EmptyString EmptyString name) ;;
  newname <- tmFreshName newname ;;
  mkVariable newname (fixNames IH);;
  mkDefinitionType name P body_syn.

Notation "'Modular' 'Fixpoint' na 'where' T_ext 'extends' T 'with' c ':=' p" := (@ModularFixpointN (fun na : nat => na) T T_ext _ _ (fun c => p)) (at level 1, T_ext at next level, T at next level, c at next level, p at next level).

Notation "'Modular' 'Fixpoint' na 'where' T_ext 'extends' T ':=' p" := (@ModularFixpoint2 (fun na : nat => na) T T_ext _ p) (at level 1, T_ext at next level, T at next level, p at next level).

Hint Extern 0 => reflexivity.

(* Definition tmkAddInjection (name : ident) T := *)
(*   T_syn <- tmQuote T ;; *)
(*   let '((args, H1), H2) := split_forall_impl  T_syn in *)
(*   mkVariable name T;;  *)
(*   let I_syn := buildImp args H1 H2 in *)
(*   I <- tmUnquoteTyped Type I_syn;; *)
(*   i <- tmUnquoteTyped I (tConst name Instance.empty)  ;; *)
(*   (tmBind (tmDefinition (append name "_inst") i) (fun _ => tmReturn tt)). *)