From Hammer Require Import Hammer.

Require Import Rbase.
Require Import Rbasic_fun.
Require Import Even.
Require Import Div2.
Require Import ArithRing.

Local Open Scope Z_scope.
Local Open Scope R_scope.

Lemma minus_neq_O : forall n i:nat, (i < n)%nat -> (n - i)%nat <> 0%nat.
Proof. hammer_hook "ArithProp" "ArithProp.minus_neq_O".
intros; red; intro.
cut (forall n m:nat, (m <= n)%nat -> (n - m)%nat = 0%nat -> n = m).
intro; assert (H2 := H1 _ _ (lt_le_weak _ _ H) H0); rewrite H2 in H;
elim (lt_irrefl _ H).
set (R := fun n m:nat => (m <= n)%nat -> (n - m)%nat = 0%nat -> n = m).
cut
((forall n m:nat, R n m) ->
forall n0 m:nat, (m <= n0)%nat -> (n0 - m)%nat = 0%nat -> n0 = m).
intro; apply H1.
apply nat_double_ind.
unfold R; intros; inversion H2; reflexivity.
unfold R; intros; simpl in H3; assumption.
unfold R; intros; simpl in H4; assert (H5 := le_S_n _ _ H3);
assert (H6 := H2 H5 H4); rewrite H6; reflexivity.
unfold R; intros; apply H1; assumption.
Qed.

Lemma le_minusni_n : forall n i:nat, (i <= n)%nat -> (n - i <= n)%nat.
Proof. hammer_hook "ArithProp" "ArithProp.le_minusni_n".
set (R := fun m n:nat => (n <= m)%nat -> (m - n <= m)%nat).
cut
((forall m n:nat, R m n) -> forall n i:nat, (i <= n)%nat -> (n - i <= n)%nat).
intro; apply H.
apply nat_double_ind.
unfold R; intros; simpl; apply le_n.
unfold R; intros; simpl; apply le_n.
unfold R; intros; simpl; apply le_trans with n.
apply H0; apply le_S_n; assumption.
apply le_n_Sn.
unfold R; intros; apply H; assumption.
Qed.

Lemma lt_minus_O_lt : forall m n:nat, (m < n)%nat -> (0 < n - m)%nat.
Proof. hammer_hook "ArithProp" "ArithProp.lt_minus_O_lt".
intros n m; pattern n, m; apply nat_double_ind;
[ intros; rewrite <- minus_n_O; assumption
| intros; elim (lt_n_O _ H)
| intros; simpl; apply H; apply lt_S_n; assumption ].
Qed.

Lemma even_odd_cor :
forall n:nat,  exists p : nat, n = (2 * p)%nat \/ n = S (2 * p).
Proof. hammer_hook "ArithProp" "ArithProp.even_odd_cor".
intro.
assert (H := even_or_odd n).
exists (div2 n).
assert (H0 := even_odd_double n).
elim H0; intros.
elim H1; intros H3 _.
elim H2; intros H4 _.
replace (2 * div2 n)%nat with (double (div2 n)).
elim H; intro.
left.
apply H3; assumption.
right.
apply H4; assumption.
unfold double;ring.
Qed.

Lemma le_double : forall m n:nat, (2 * m <= 2 * n)%nat -> (m <= n)%nat.
Proof. hammer_hook "ArithProp" "ArithProp.le_double".
intros; apply INR_le.
assert (H1 := le_INR _ _ H).
do 2 rewrite mult_INR in H1.
apply Rmult_le_reg_l with (INR 2).
replace (INR 2) with 2; [ prove_sup0 | reflexivity ].
assumption.
Qed.

Lemma euclidian_division :
forall x y:R,
y <> 0 ->
exists k : Z, (exists r : R, x = IZR k * y + r /\ 0 <= r < Rabs y).
Proof. hammer_hook "ArithProp" "ArithProp.euclidian_division".
intros.
set
(k0 :=
match Rcase_abs y with
| left _ => (1 - up (x / - y))%Z
| right _ => (up (x / y) - 1)%Z
end).
exists k0.
exists (x - IZR k0 * y).
split.
ring.
unfold k0; case (Rcase_abs y) as [Hlt|Hge].
assert (H0 := archimed (x / - y)); rewrite <- Z_R_minus; simpl;
unfold Rminus.
replace (- ((1 + - IZR (up (x / - y))) * y)) with
((IZR (up (x / - y)) - 1) * y); [ idtac | ring ].
split.
apply Rmult_le_reg_l with (/ - y).
apply Rinv_0_lt_compat; apply Ropp_0_gt_lt_contravar; exact Hlt.
rewrite Rmult_0_r; rewrite (Rmult_comm (/ - y)); rewrite Rmult_plus_distr_r;
rewrite <- Ropp_inv_permute; [ idtac | assumption ].
rewrite Rmult_assoc; repeat rewrite Ropp_mult_distr_r_reverse;
rewrite <- Rinv_r_sym; [ rewrite Rmult_1_r | assumption ].
apply Rplus_le_reg_l with (IZR (up (x / - y)) - x / - y).
rewrite Rplus_0_r; unfold Rdiv; pattern (/ - y) at 4;
rewrite <- Ropp_inv_permute; [ idtac | assumption ].
replace
(IZR (up (x * / - y)) - x * - / y +
(- (x * / y) + - (IZR (up (x * / - y)) - 1))) with 1;
[ idtac | ring ].
elim H0; intros _ H1; unfold Rdiv in H1; exact H1.
rewrite (Rabs_left _ Hlt); apply Rmult_lt_reg_l with (/ - y).
apply Rinv_0_lt_compat; apply Ropp_0_gt_lt_contravar; exact Hlt.
rewrite <- Rinv_l_sym.
rewrite (Rmult_comm (/ - y)); rewrite Rmult_plus_distr_r;
rewrite <- Ropp_inv_permute; [ idtac | assumption ].
rewrite Rmult_assoc; repeat rewrite Ropp_mult_distr_r_reverse;
rewrite <- Rinv_r_sym; [ rewrite Rmult_1_r | assumption ];
apply Rplus_lt_reg_l with (IZR (up (x / - y)) - 1).
replace (IZR (up (x / - y)) - 1 + 1) with (IZR (up (x / - y)));
[ idtac | ring ].
replace (IZR (up (x / - y)) - 1 + (- (x * / y) + - (IZR (up (x / - y)) - 1)))
with (- (x * / y)); [ idtac | ring ].
rewrite <- Ropp_mult_distr_r_reverse; rewrite (Ropp_inv_permute _ H); elim H0;
unfold Rdiv; intros H1 _; exact H1.
apply Ropp_neq_0_compat; assumption.
assert (H0 := archimed (x / y)); rewrite <- Z_R_minus; simpl;
cut (0 < y).
intro; unfold Rminus;
replace (- ((IZR (up (x / y)) + -(1)) * y)) with ((1 - IZR (up (x / y))) * y);
[ idtac | ring ].
split.
apply Rmult_le_reg_l with (/ y).
apply Rinv_0_lt_compat; assumption.
rewrite Rmult_0_r; rewrite (Rmult_comm (/ y)); rewrite Rmult_plus_distr_r;
rewrite Rmult_assoc; rewrite <- Rinv_r_sym;
[ rewrite Rmult_1_r | assumption ];
apply Rplus_le_reg_l with (IZR (up (x / y)) - x / y);
rewrite Rplus_0_r; unfold Rdiv;
replace
(IZR (up (x * / y)) - x * / y + (x * / y + (1 - IZR (up (x * / y))))) with
1; [ idtac | ring ]; elim H0; intros _ H2; unfold Rdiv in H2;
exact H2.
rewrite (Rabs_right _ Hge); apply Rmult_lt_reg_l with (/ y).
apply Rinv_0_lt_compat; assumption.
rewrite <- (Rinv_l_sym _ H); rewrite (Rmult_comm (/ y));
rewrite Rmult_plus_distr_r; rewrite Rmult_assoc; rewrite <- Rinv_r_sym;
[ rewrite Rmult_1_r | assumption ];
apply Rplus_lt_reg_l with (IZR (up (x / y)) - 1);
replace (IZR (up (x / y)) - 1 + 1) with (IZR (up (x / y)));
[ idtac | ring ];
replace (IZR (up (x / y)) - 1 + (x * / y + (1 - IZR (up (x / y))))) with
(x * / y); [ idtac | ring ]; elim H0; unfold Rdiv;
intros H2 _; exact H2.
destruct (total_order_T 0 y) as [[Hlt|Heq]|Hgt].
assumption.
elim H; symmetry ; exact Heq.
apply Rge_le in Hge; elim (Rlt_irrefl _ (Rle_lt_trans _ _ _ Hge Hgt)).
Qed.

Lemma tech8 : forall n i:nat, (n <= S n + i)%nat.
Proof. hammer_hook "ArithProp" "ArithProp.tech8".
intros; induction  i as [| i Hreci].
replace (S n + 0)%nat with (S n); [ apply le_n_Sn | ring ].
replace (S n + S i)%nat with (S (S n + i)).
apply le_S; assumption.
apply INR_eq; rewrite S_INR; do 2 rewrite plus_INR; do 2 rewrite S_INR; ring.
Qed.


Require
  MathClasses.theory.setoids.
Require Import
  Coq.Classes.Morphisms MathClasses.interfaces.abstract_algebra.

Section semigroup_props.
Context `{SemiGroup G}.

Global Instance sg_op_mor_1: ∀ x, Setoid_Morphism (x &) | 0.
Proof. split; try apply _. Qed.

Global Instance sg_op_mor_2: ∀ x, Setoid_Morphism (& x) | 0.
Proof. split; try apply _. solve_proper. Qed.
End semigroup_props.

(**
 * This tactic simplifies group expressions by first pushing group morphisms and
 * negations down to the leaves, then cancelling inverses and removing units.
 *)
Ltac group_simplify :=
  rewrite_strat
    (try bottomup (hints group_simplify));
    (try bottomup (choice (hints group_cancellation) <- associativity)).

Ltac group := group_simplify; easy.

Set Warnings "-unsupported-attributes". (* FIXME: remove when minimal Coq version is enough *)

#[global]
Hint Rewrite @associativity using apply _: group_simplify.
#[global]
Hint Rewrite @left_identity @right_identity @left_inverse @right_inverse using apply _: group_cancellation.

Set Warnings "+unsupported-attributes".

Section group_props.
Context `{Group G}.
Lemma negate_sg_op_distr_flip `{Group A}: ∀ a b, -(a & b) = -b & -a.
Proof.
intros.
  setoid_replace (-b & -a) with (-(a & b) & (a & b) & (-b & -a)) by group;
  rewrite <- associativity;
  setoid_replace (a & b & (-b & -a)) with mon_unit;
  group.
Qed.

Hint Rewrite @negate_sg_op_distr_flip using apply _ : group_simplify.

Global Instance negate_involutive: Involutive (-).
Proof.
  intros x;
  setoid_replace x with (- - x & (- x & x)) at 2 by group;
  setoid_replace (-x & x) with mon_unit by group;
  group.
Qed.

Hint Rewrite @negate_involutive using apply _ : group_cancellation.

Global Instance: Injective (-).
Proof.
  repeat (split; try apply _).
  intros x y E.
  now rewrite <-(involutive x), <-(involutive y), E.
Qed.

Lemma negate_mon_unit : -mon_unit = mon_unit.
Proof.
  setoid_replace mon_unit with (-mon_unit & mon_unit); group.
Qed.

Hint Rewrite @negate_mon_unit using apply _ : group_cancellation.

Global Instance: ∀ z : G, LeftCancellation (&) z.
Proof.
  intros z x y E;
  setoid_replace x with (-z & (z & x)) by group;
  rewrite E;
  group.
Qed.

Global Instance: ∀ z : G, RightCancellation (&) z.
Proof.
  intros z x y E;
  setoid_replace x with (x & z & -z) by group;
  rewrite E;
  group.
Qed.

Lemma negate_sg_op_distr `{!AbGroup G} x y: -(x & y) = -x & -y.
Proof. group_simplify. apply commutativity. Qed.

End group_props.

Section groupmor_props.
  Context `{Group A} `{Group B} {f : A → B} `{!Monoid_Morphism f}.

  Lemma preserves_negate x : f (-x) = -f x.
  Proof.
    apply (left_cancellation (&) (f x));
    rewrite <- preserves_sg_op;
    group_simplify;
    apply preserves_mon_unit.
  Qed.

  Hint Rewrite @preserves_sg_op @preserves_negate @preserves_mon_unit using apply _ : group_simplify.
End groupmor_props.

#[global]
Instance semigroup_morphism_proper {A B eA eB opA opB} :
  Proper ((=) ==> iff) (@SemiGroup_Morphism A B eA eB opA opB) | 1.
Proof.
  assert (∀ (f g : A → B), g = f → SemiGroup_Morphism f → SemiGroup_Morphism g) as P.
   intros f g E [? ? ? ?].
   split; try apply _.
    eapply setoids.morphism_proper; eauto.
   intros x y. now rewrite (E (x & y)), (E x), (E y).
  intros f g ?; split; intros Mor.
   apply P with f. destruct Mor. now symmetry. apply _.
  now apply P with g.
Qed.

#[global]
Instance monoid_morphism_proper {A B eA eB opA uA opB uB} :
  Proper ((=) ==> iff) (@Monoid_Morphism A B eA eB opA uA opB uB) | 1.
Proof.
  assert (∀ (f g : A → B), g = f → Monoid_Morphism f → Monoid_Morphism g) as P.
   intros f g E [? ? ? ?].
   split; try apply _.
    eapply semigroup_morphism_proper; eauto.
   now rewrite (E mon_unit mon_unit).
  intros f g ?; split; intros Mor.
   apply P with f. destruct Mor. now symmetry. apply _.
  now apply P with g.
Qed.

Section from_another_sg.
  Context `{SemiGroup A} `{Setoid B}
   `{Bop : SgOp B} (f : B → A) `{!Injective f} (op_correct : ∀ x y, f (x & y) = f x & f y).

  Instance: Setoid_Morphism f := injective_mor f.
  Instance: Proper ((=) ==> (=) ==> (=)) Bop.
  Proof. intros ? ? E1 ? ? E2. apply (injective f). rewrite 2!op_correct. apply sg_op_proper; now apply sm_proper. Qed.

  Lemma projected_sg: SemiGroup B.
  Proof.
    split; try apply _.
    repeat intro; apply (injective f). now rewrite !op_correct, associativity.
  Qed.
End from_another_sg.

Section from_another_com_sg.
  Context `{CommutativeSemiGroup A} `{Setoid B}
   `{Bop : SgOp B} (f : B → A) `{!Injective f} (op_correct : ∀ x y, f (x & y) = f x & f y).

  Lemma projected_com_sg: CommutativeSemiGroup B.
  Proof.
    split. now apply (projected_sg f).
    repeat intro; apply (injective f). now rewrite !op_correct, commutativity.
  Qed.
End from_another_com_sg.

Section from_another_monoid.
  Context `{Monoid A} `{Setoid B}
   `{Bop : SgOp B} `{Bunit : MonUnit B} (f : B → A) `{!Injective f}
   (op_correct : ∀ x y, f (x & y) = f x & f y) (unit_correct : f mon_unit = mon_unit).

  Lemma projected_monoid: Monoid B.
  Proof.
    split. now apply (projected_sg f).
     repeat intro; apply (injective f). now rewrite op_correct, unit_correct, left_identity.
    repeat intro; apply (injective f). now rewrite op_correct, unit_correct, right_identity.
  Qed.
End from_another_monoid.

Section from_another_com_monoid.
  Context `{CommutativeMonoid A} `{Setoid B}
   `{Bop : SgOp B} `{Bunit : MonUnit B} (f : B → A) `{!Injective f}
   (op_correct : ∀ x y, f (x & y) = f x & f y) (unit_correct : f mon_unit = mon_unit).

  Lemma projected_com_monoid: CommutativeMonoid B.
  Proof.
    split. now apply (projected_monoid f).
    repeat intro; apply (injective f). now rewrite op_correct, commutativity.
  Qed.
End from_another_com_monoid.

Section from_another_group.
  Context `{Group A} `{Setoid B}
   `{Bop : SgOp B} `{Bunit : MonUnit B} `{Bnegate : Negate B} (f : B → A) `{!Injective f}
   (op_correct : ∀ x y, f (x & y) = f x & f y) (unit_correct : f mon_unit = mon_unit)
   (negate_correct : ∀ x, f (-x) = -f x).

  Instance: Setoid_Morphism f := injective_mor f.
  Instance: Setoid_Morphism Bnegate.
  Proof. split; try apply _. intros ? ? E1. apply (injective f). rewrite 2!negate_correct. now do 2 apply sm_proper. Qed.

  Lemma projected_group: Group B.
  Proof.
    split; try apply _. now apply (projected_monoid f).
     repeat intro; apply (injective f). now rewrite op_correct, negate_correct, unit_correct, left_inverse.
    repeat intro; apply (injective f). now rewrite op_correct, negate_correct, unit_correct, right_inverse.
  Qed.
End from_another_group.

Section from_another_ab_group.
  Context `{AbGroup A} `{Setoid B}
   `{Bop : SgOp B} `{Bunit : MonUnit B} `{Bnegate : Negate B} (f : B → A) `{!Injective f}
   (op_correct : ∀ x y, f (x & y) = f x & f y) (unit_correct : f mon_unit = mon_unit)
   (negate_correct : ∀ x, f (-x) = -f x).

  Lemma projected_ab_group: AbGroup B.
  Proof.
    split; try apply _. now apply (projected_group f).
    repeat intro; apply (injective f). now rewrite op_correct, commutativity.
  Qed.
End from_another_ab_group.


Require Import MapProofs.Common.
Set Bullet Behavior "Strict Subproofs".
Require Import MapProofs.Bounds.
Require Import MapProofs.Tactics.
Require Import MapProofs.UnionIntersectDifferenceProofs.
Require Import MapProofs.ToListProofs.
Require Import Coq.Classes.Morphisms.

Section WF.
Context {e : Type} {a : Type} {HEq : Eq_ e} {HOrd : Ord e} {HEqLaws : EqLaws e}  {HOrdLaws : OrdLaws e}.

(** ** Verification of [submap] *)
Lemma submap'_spec:
  forall (m1: Map e a) (m2: Map e a) lb ub f,
  Bounded m1 lb ub ->
  Bounded m2 lb ub ->
  submap' f m1 m2 = true <->
  (forall i value, sem m1 i = Some value -> exists v, (sem m2 i = Some v /\ f value v = true)).
Proof.
  intros ????? HB1 HB2.
  revert dependent m2.
  induction HB1; intros; simpl; subst.
  * intuition. discriminate H0.
  * destruct m2 eqn:Hs0.
    - rewrite <- Hs0 in *.
      clear s e0 a0 m1 m3 Hs0.
      eapply splitLookup_Desc; [solve_Bounded e|].
      intros sr1 b sr2 HBsr1 HBsr2 Hb Hsem.
      destruct HB2.
      + simpl. split; intros. simpl in Hb. subst. discriminate H1. specialize (H1 x v).
        assert (sem s1 x = None). { eapply sem_outside_above. apply HB1_1. unfold isUB. order e. }
        rewrite H3 in H1. simpl in H1. rewrite Eq_Reflexive in H1. simpl in H1.
        assert (exists v0 : a, None = Some v0 /\ f v v0 = true) by  (apply H1; reflexivity).
        destruct H4. destruct H4. discriminate H4.
      + destruct b. rewrite andb_true_iff. rewrite andb_true_iff. rewrite IHHB1_1 by eassumption.
        rewrite IHHB1_2 by eassumption. split; intro.
        -- destruct H6. destruct H7. intros i value Hi. rewrite Hsem. destruct (sem s1 i) eqn : ?.
            ** apply H7 in Heqo. destruct Heqo. destruct H9. exists x1.
               destruct (i == x) eqn : ?. solve_Bounds e. rewrite H9. simpl. simpl in Hi. split. reflexivity.
              inversion Hi. subst. apply H10.
            ** simpl in Hi. destruct (i == x) eqn : ?. simpl in Hi. exists a0. split.
              eapply sem_resp_eq in Heqb. rewrite <- Hb in Heqb. assumption. inversion Hi; subst.
              assumption. simpl in Hi. apply H8 in Hi. destruct Hi. exists x1.
              assert (sem sr1 i = None). { eapply sem_outside_above. apply HBsr1.
              destruct H9. apply (sem_inside HBsr2) in H9. destruct H9. solve_Bounds e. }
              rewrite H10. simpl. apply H9.
       -- split. specialize (H6 x v). assert (sem s1 x = None). { eapply sem_outside_above.
          apply HB1_1. unfold isUB. order e. } rewrite H7 in H6. simpl in H6. rewrite Eq_Reflexive in H6.
          simpl in H6.
          assert (exists v1 : a, sem s0 x ||| SomeIf (_GHC.Base.==_ x x0) v0 ||| sem s3 x = Some v1
          /\ f v v1 = true) by (apply H6; reflexivity). destruct H8. simpl in Hb. rewrite <- Hb in H8.
          destruct H8. inversion H8; subst. assumption.
          split.
          ** intros. specialize (H6 i value). rewrite H7 in H6. simpl in H6.
             assert (exists v : a, sem s0 i ||| SomeIf (_GHC.Base.==_ i x0) v0 ||| sem s3 i = Some v
             /\ f value v = true) by (apply H6; reflexivity). destruct H8. destruct H8.
            specialize (Hsem i). simpl in Hsem. rewrite H8 in Hsem. destruct (i==x) eqn : ?.
            ++ solve_Bounds e.
            ++ assert (sem sr2 i = None). { eapply (sem_inside HB1_1) in H7. destruct H7.
              unfold isUB in H10. eapply (sem_outside_below). apply HBsr2. solve_Bounds e. }
              rewrite H10 in Hsem. rewrite oro_None_r in Hsem. exists x1. split. symmetry.
              apply Hsem. apply H9.
          ** intros. specialize (H6 i value). rewrite H7 in H6.
              assert (sem s1 i = None). eapply sem_outside_above. apply HB1_1. eapply (sem_inside HB1_2) in H7.
              destruct H7. solve_Bounds e. rewrite H8 in H6. simpl in H6. destruct (i == x) eqn : ?.
              ++ solve_Bounds e.
              ++ simpl in H6. assert ( exists v : a, sem s0 i |||
              SomeIf (_GHC.Base.==_ i x0) v0 ||| sem s3 i = Some v /\ f value v = true) by (apply H6; reflexivity).
             destruct H9. specialize (Hsem i). simpl in Hsem. destruct H9. rewrite H9 in Hsem.
              rewrite Heqb in Hsem. assert (sem sr1 i = None). { eapply (sem_inside HB1_2) in H7.
              destruct H7. unfold isLB in H7. eapply sem_outside_above. apply HBsr1. solve_Bounds e. }
              rewrite H11 in Hsem. simpl in Hsem. exists x1. split. symmetry. apply Hsem. apply H10.
          -- split; intros. discriminate H6. specialize (H6 x). assert (sem s1 x = None). {
              eapply sem_outside_above. apply HB1_1. unfold isUB.  order e. } rewrite H7 in H6. simpl in H6.
              specialize (H6 v). rewrite Eq_Reflexive in H6. simpl in H6. destruct H6. reflexivity.
              simpl in Hb. rewrite <- Hb in H6. destruct H6. discriminate H6.
    - intuition.
      + discriminate H1.
      + subst. simpl in H1. specialize (H1 x v). assert (sem s1 x = None). { eapply sem_outside_above.
        apply HB1_1. unfold isUB. order e. } rewrite H3 in H1. simpl in H1. rewrite Eq_Reflexive in H1.
        simpl in H1. destruct H1. reflexivity. destruct H1. discriminate H1.
Qed.

Lemma submap_size :
  forall (m1: Map e a) (m2: Map e a) lb ub,
  Bounded m1 lb ub ->
  Bounded m2 lb ub ->
  (forall i value, sem m1 i = Some value ->  exists v, sem m2 i = Some v) ->
  (size m1 <= size m2)%Z.
Proof.
  intros ???? HB1 HB2 Hsubmap.
  revert dependent m2.
  induction HB1; intros; simpl; subst.
  * simpl. solve_size.
  * assert (exists v, sem m2 x = Some v). { specialize (Hsubmap x v). simpl in Hsubmap.
    assert (sem s1 x = None). { eapply sem_outside_above. apply HB1_1. unfold isUB. order e. }
    rewrite H1 in Hsubmap. simpl in Hsubmap. rewrite Eq_Reflexive in Hsubmap. simpl in Hsubmap.
    destruct Hsubmap. reflexivity. exists x0. assumption. }
    assert (size m2 = let '(sl,sr) := split x m2 in 1 + size sl + size sr)%Z.
    { eapply split_Desc; [eassumption|]. intros. destruct H1. rewrite H1 in H5. simpl in H5. lia. }
    rewrite H3.
    eapply split_Desc; [eassumption|]. intros.
    assert (size s1 <= size s0)%Z.
    { apply IHHB1_1; try assumption.
      intros i v1 Hi.
      specialize (Hsubmap i). simpl in Hsubmap.
      rewrite Hi in Hsubmap. simpl in Hsubmap.
      specialize (Hsubmap v1). destruct Hsubmap. reflexivity.
      specialize (H7 i). destruct (i == x) eqn : ?.
      { solve_Bounds e. }
      { assert (sem s3 i = None). { assert (i < x = true) by solve_Bounds e. eapply sem_outside_below.
        apply H5. unfold isLB. order e. }
        rewrite H9 in H7. rewrite oro_None_r  in H7. exists x0. rewrite <- H7. assumption. }
   }
    assert (size s2 <= size s3)%Z.
    { apply IHHB1_2; try assumption.
      intros i v1 Hi.
      specialize (Hsubmap i). simpl in Hsubmap.
      rewrite Hi in Hsubmap. simpl in Hsubmap.
      specialize (Hsubmap v1).
      assert (sem s1 i = None). { eapply sem_outside_above. apply HB1_1. solve_Bounds e. }
      rewrite H9 in Hsubmap. simpl in Hsubmap.
      assert (i == x = false) by solve_Bounds e. rewrite H10 in Hsubmap. simpl in Hsubmap.
      destruct Hsubmap. reflexivity. specialize (H7 i). rewrite H11 in H7. rewrite H10 in H7.
     assert (sem s0 i = None). { assert (x < i = true) by solve_Bounds e. eapply sem_outside_above.
      apply H4. solve_Bounds e. }
      rewrite H12 in H7. simpl in H7. exists x0. symmetry. apply H7. }
    lia.
Qed.

Lemma isSubmapOfBy_spec:
  forall (m1: Map e a) (m2: Map e a) f lb ub,
  Bounded m1 lb ub ->
  Bounded m2 lb ub ->
  isSubmapOfBy f m1 m2 = true <-> (forall i value, sem m1 i = Some value -> exists v, sem m2 i =
  Some v /\ f value v = true).
Proof.
  intros. unfold isSubmapOfBy. split; intros.
  - eapply submap'_spec. apply H. apply H0. rewrite andb_true_iff in H1. destruct H1. apply H3.
    apply H2.
  - apply andb_true_iff. split. unfold op_zlze__. unfold Ord_Integer___. unfold op_zlze____.
    rewrite size_size. rewrite Z.leb_le. eapply submap_size. apply H. apply H0. intros.
    specialize (H1 i value). apply H1 in H2. destruct H2. destruct H2. exists x. apply H2.
    eapply submap'_spec. apply H. apply H0. apply H1.
Qed.

Lemma isSubmapOf_spec:
  forall `{Eq_ a} m1 m2 lb ub,
  Bounded m1 lb ub ->
  Bounded m2 lb ub ->
  isSubmapOf m1 m2 = true <-> (forall i value, sem m1 i = Some value -> exists v, sem m2 i = Some v /\
  value == v = true).
Proof.
  intros. eapply isSubmapOfBy_spec. apply H0. apply H1.
Qed.

(** ** Verification of [isProperSubmapOfBy] *)
Lemma isProperSubmapOfBy_spec:
  forall (m1: Map e a) (m2: Map e a) lb ub f,
  Bounded m1 lb ub ->
  Bounded m2 lb ub ->
    isProperSubmapOfBy f m1 m2 = true <-> isSubmapOfBy f m1 m2 = true /\ size m1 < size m2 = true.
Proof.
  intros. unfold isProperSubmapOfBy. unfold isSubmapOfBy. repeat (rewrite andb_true_iff in *). split; intros.
   - destruct H1. split. split. order Int. assumption. order Int.
   - destruct H1. destruct H1. split. order Int. apply H3.
Qed.

(** ** Veriticcation of [isProperSubmpaOf] *)
Lemma isProperSubmapOf_spec:
  forall `{EqLaws a} (m1: Map e a) (m2: Map e a) lb ub,
  Bounded m1 lb ub ->
  Bounded m2 lb ub ->
  isProperSubmapOf m1 m2 = true <-> isSubmapOf m1 m2 = true /\ size m1 < size m2 = true.
Proof.
  intros. unfold isProperSubmapOf. unfold isSubmapOf. eapply isProperSubmapOfBy_spec; eassumption.
Qed.

(** ** Verification of [filter] *)

(**
For filter we need two lemmas: We need to know that [filter P s] is
well-formed even if P does not respect equality (this is
required by the [FSetInterface]). But to prove something about its
semantics, we need to assume that [P] respects equality.
*)

Lemma filterWithKey_Bounded:
  forall (P : e -> a -> bool) map lb ub,
  Bounded map lb ub ->
  Bounded (Internal.filterWithKey P map) lb ub.
Proof.
  intros.
  induction H.
  * simpl. solve_Bounded e.
  * simpl.
    destruct (P x v) eqn:HPx.
    - destruct_ptrEq.
      + solve_Bounded e.
      + applyDesc e (@link_Desc e a).
    - applyDesc e (@link2_Desc e a).
Qed.

Lemma filterWithKey_Desc:
  forall (P : e -> a -> bool) map lb ub,
  Bounded map lb ub ->
  Proper ((fun i j : e => _GHC.Base.==_ i j = true) ==> eq) P ->
  Desc' (Internal.filterWithKey P map) lb ub (fun x => match sem map x with
                                                      | Some y => if P x y then Some y else None
                                                      | None => None
                                                      end).
Proof.
  intros.
  induction H.
  * simpl. solve_Desc e. reflexivity.
  * simpl.
    applyDesc e IHBounded1.
    applyDesc e IHBounded2.
    destruct (P x v) eqn:HPx.
    - destruct_ptrEq.
      + solve_Desc e. f_solver e. assert (P i v = true). { erewrite H0. apply HPx. apply Heqb. }
        rewrite Heqb0 in H4. inversion H4.
      + applyDesc e (@link_Desc e a).
        solve_Desc e. f_solver e. assert (P i a0 = true). { erewrite H0. apply HPx. apply Heqb. }
        rewrite Heqb0 in H4. inversion H4.
    - applyDesc e (@link2_Desc e a).
      solve_Desc e. f_solver e. assert (P x v = true). { erewrite H0. apply Heqb0. apply Eq_Symmetric.
      apply Heqb. } rewrite HPx in H6. inversion H6.
Qed.

Lemma filterWithKey_WF :
  forall f (s : Map e a),
    Proper ((fun i j : e => _GHC.Base.==_ i j = true) ==> eq) f ->
    WF s -> WF (filterWithKey f s).
Proof.
  intros. eapply Desc'_WF.
  apply filterWithKey_Desc; assumption.
Qed.

(*This requires no conditions on the function P*)
Lemma filter_Desc:
  forall  (P : a -> bool) map lb ub,
  Bounded map lb ub ->
  Desc' (Internal.filter P map) lb ub (fun x => match sem map x with
                                                |Some y => if P y then Some y else None
                                                |None => None
                                                end).
Proof.
  intros. eapply filterWithKey_Desc. apply H. f_solver e.
Qed.

Lemma filter_WF :
  forall f (s : Map e a),
    WF s -> WF (filter f s).
Proof.
  intros. eapply Desc'_WF.
  apply filter_Desc; assumption.
Qed.

(** ** Verification of [partition] *)

Lemma partitionWithKey_Bounded:
  forall p map lb ub,
  Bounded map lb ub ->
  forall (P : (Map e a * Map e a) -> Prop),
  (forall m1 m2, Bounded m1 lb ub /\ Bounded m2 lb ub -> P (m1, m2)) ->
  P (partitionWithKey p map).
Proof.
  intros ???? HB.
  induction HB.
  * intros X HX; apply HX; clear X HX; split; solve_Bounded e.
  * simpl.
    apply IHHB1; intros sl1 sl2 [HDsl1 HDsl2]; clear IHHB1.
    apply IHHB2; intros sr1 sr2 [HDsr1 HDsr2]; clear IHHB2.
    destruct (p x) eqn:?.
    - intros X HX; apply HX; clear X HX; split.
      + destruct_ptrEq.
        -- solve_Bounded e.
        -- applyDesc e (@link_Desc e a).
      + applyDesc e (@link2_Desc e a).
    - intros X HX; apply HX; clear X HX; split.
      + applyDesc e (@link2_Desc e a).
      + destruct_ptrEq.
        -- solve_Bounded e.
        -- applyDesc e (@link_Desc e a).
Qed.

Lemma partitionWithKey_spec:
  forall (p : e -> a -> bool) map lb ub,
  Bounded map lb ub ->
  Proper ((fun i j : e => i == j = true) ==> eq) p ->
  forall (P : (Map e a * Map e a) -> Prop),
  (forall m1 m2,
    Desc' m1 lb ub (fun i => match sem map i with
                             | Some y => if p i y then Some y else None
                             | None => None
                              end)  /\
    Desc' m2 lb ub (fun i => match sem map i with
                             | Some y => if p i y then None else Some y
                             | None => None
                             end) ->
    P (m1, m2)) ->
  P (partitionWithKey p map).
Proof.
  intros ???? HB HProper.
  induction HB.
  * intros X HX; apply HX; clear X HX; split; solve_Desc e; f_solver e.
  * simpl.
    apply IHHB1; intros sl1 sl2 [HDsl1 HDsl2]; clear IHHB1.
    applyDesc e HDsl1; clear HDsl1.
    applyDesc e HDsl2; clear HDsl2.
    apply IHHB2; intros sr1 sr2 [HDsr1 HDsr2]; clear IHHB2.
    applyDesc e HDsr1; clear HDsr1.
    applyDesc e HDsr2; clear HDsr2.
    destruct (p x) eqn:?.
    - intros X HX; apply HX; clear X HX; split.
      + destruct_ptrEq.
        -- solve_Desc e. f_solver e. assert (p i v = p x v). apply equal_f. apply HProper. assumption.
            rewrite Heqb in H1. rewrite Heqb1 in H1. inversion H1.
        -- applyDesc e (@link_Desc e a). solve_Desc e. f_solver e. assert (p i a0 = p x a0). apply equal_f.
           apply HProper. assumption. rewrite Heqb1 in H1. rewrite Heqb in H1. inversion H1.
      + applyDesc e (@link2_Desc e a). solve_Desc e. f_solver e. assert (p i v = p x v). apply equal_f.
        apply HProper. assumption. rewrite Heqb in H3. rewrite Heqb1 in H3. inversion H3.
    - intros X HX; apply HX; clear X HX; split.
      + applyDesc e (@link2_Desc e a). solve_Desc e. f_solver e. assert (p i v = p x v). apply equal_f.
        apply HProper. assumption. rewrite Heqb in H3. rewrite Heqb1 in H3. inversion H3.
      + destruct_ptrEq.
        -- solve_Desc e. f_solver e. assert (p x v = p i v). apply equal_f. apply HProper. order e.
           rewrite Heqb in H1. rewrite Heqb1 in H1. inversion H1.
        -- applyDesc e (@link_Desc e a). solve_Desc e. f_solver e. assert (p x a0 = p i a0). apply equal_f.
           apply HProper. order e. rewrite Heqb in H1. rewrite Heqb1 in H1. inversion H1.
Qed.

Lemma partition_spec:
  forall (p : a -> bool) map lb ub,
  Bounded map lb ub ->
  forall (P : (Map e a * Map e a) -> Prop),
  (forall m1 m2,
    Desc' m1 lb ub (fun i => match sem map i with
                             | Some y => if p y then Some y else None
                             | None => None
                              end)  /\
    Desc' m2 lb ub (fun i => match sem map i with
                             | Some y => if p y then None else Some y
                             | None => None
                             end) ->
    P (m1, m2)) ->
  P (partition p map).
Proof.
  intros. unfold partition. eapply partitionWithKey_spec. apply H. f_solver e. apply H0.
Qed.

(** ** Verification of [take] *)
(*This is exactly the same as SetProofs, since the take function works the exact same way*)
Definition takeGo : Int -> Map e a -> Map e a.
Proof.
  let rhs := eval unfold take in (@take e a) in
  match rhs with fun n s => if _ then _ else ?go _ _  => exact go end.
Defined.

Lemma take_neg: forall a n (l : list a), (n <= 0)%Z -> List.take n l = nil.
Proof.
  intros. destruct l; simpl; destruct (Z.leb_spec n 0); try lia; try reflexivity.
Qed.

Lemma take_all:
  forall a n (xs : list a), (Z.of_nat (length xs) <= n)%Z -> List.take n xs = xs.
Proof.
  intros. revert n H.
  induction xs; intros n Hall.
  * simpl. destruct (Z.leb_spec n 0); reflexivity.
  * simpl.
    simpl length in Hall. rewrite Nat2Z.inj_succ, <- Z.add_1_l in Hall.
    rewrite IHxs by lia.
    destruct (Z.leb_spec n 0); [lia|reflexivity].
Qed.

Lemma take_app_cons:
  forall a n (l1 : list a) (x : a) (l2 : list a),
  List.take n (l1 ++ x :: l2) = match (n ?= Z.of_nat (length l1))%Z with
    | Lt => List.take n l1
    | Eq => l1
    | Gt => l1 ++ (x :: nil) ++ List.take (n - Z.of_nat (length l1) - 1)%Z l2
  end.
Proof.
  intros. revert n.
  induction l1; intros n.
  * simpl.
    rewrite Z.sub_0_r.
    destruct (Z.leb_spec n 0), (Z.compare_spec n 0)%Z; try reflexivity; lia.
  * cbn -[Z.add Z.of_nat].
    rewrite IHl1. clear IHl1.
    rewrite Nat2Z.inj_succ, <- Z.add_1_l.
    replace (n - (1 + Z.of_nat (length l1)) - 1)%Z with (n - 2 - Z.of_nat (length l1))%Z by lia.
    replace (n - 1 - (Z.of_nat (length l1)) - 1)%Z with (n - 2 - Z.of_nat (length l1))%Z by lia.
    destruct (Z.leb_spec n 0),
             (Z.compare_spec (n - 1) (Z.of_nat (length l1)))%Z,
             (Z.compare_spec n (1 + Z.of_nat (length l1)))%Z; try reflexivity; lia.
Qed.

Lemma takeGo_spec :
  forall n map lb ub,
  Bounded map lb ub ->
  forall (P : Map e a -> Prop),
  (forall m',
    Bounded m' lb ub /\
    toList m' = List.take n (toList map) ->
    P m') ->
  P (takeGo n map).
Proof.
  intros ???? HD. revert n.
  induction HD; intros n.
  * intros X HX; apply HX. split.
    + simpl. destruct_match; solve_Bounded e.
    + simpl. do 2 destruct_match; reflexivity.
  * simpl.
    intros X HX; apply HX.
    rewrite toList_Bin.
    unfold op_zlze__ , Ord_Integer___, op_zlze____.
    unfold compare , Ord_Integer___, compare__.
    rewrite size_size.
    apply IHHD1. intros s1' [HBs1' Hs1']. clear IHHD1.
    apply IHHD2. intros s2' [HBs2' Hs2']. clear IHHD2.
    destruct (Z.leb_spec n 0).
    + rewrite take_neg by assumption. split.
      - solve_Bounded e.
      - reflexivity.
    + simpl app.
      rewrite take_app_cons.
      erewrite <- size_spec by eassumption.
      destruct (Z.compare_spec n (size s1)).
      - split.
        ** eapply Bounded_relax_ub; eassumption.
        ** reflexivity.
      - split.
        ** eapply Bounded_relax_ub; eassumption.
        ** assumption.
      - split.
        ** applyDesc e (@link_Desc e a).
        ** erewrite toList_link by solve_Precondition e.
           rewrite Hs2'.
           reflexivity.
Qed.

Lemma toList_take:
  forall n map lb ub,
  Bounded map lb ub ->
  forall (P : Map e a -> Prop),
  (forall m',
    Bounded m' lb ub /\
    toList m' = List.take n (toList map) ->
    P m') ->
  P (take n map).
Proof.
  intros. apply H0.
  unfold take. fold takeGo.
  unfold op_zgze__ , Ord_Integer___, op_zgze____.
  rewrite size_size.
  destruct (Z.leb_spec (size map) n).
  * split; [assumption|].
    rewrite take_all. reflexivity.
    erewrite <- size_spec by eassumption.
    assumption.
  * eapply takeGo_spec; [solve_Precondition e..|].
    intros. assumption.
Qed.

(** ** Verification of [drop] *)

Definition dropGo : Int -> Map e a -> Map e a.
Proof.
  let rhs := eval unfold drop in (@drop e a) in
  match rhs with fun n s => if _ then _ else ?go _ _  => exact go end.
Defined.

Lemma drop_neg: forall a n (l : list a), (n <= 0)%Z -> List.drop n l = l.
Proof.
  intros. destruct l; simpl; destruct (Z.leb_spec n 0); try lia; try reflexivity.
Qed.

Lemma drop_all:
  forall a n (xs : list a), (Z.of_nat (length xs) <= n)%Z -> List.drop n xs = nil.
Proof.
  intros. revert n H.
  induction xs; intros n Hall.
  * simpl. destruct (Z.leb_spec n 0); reflexivity.
  * simpl.
    simpl length in Hall. rewrite Nat2Z.inj_succ, <- Z.add_1_l in Hall.
    rewrite IHxs by lia.
    destruct (Z.leb_spec n 0); [lia|reflexivity].
Qed.

Lemma drop_app_cons:
  forall a n (l1 : list a) (x : a) (l2 : list a),
  List.drop n (l1 ++ x :: l2) = match (n ?= Z.of_nat (length l1))%Z with
    | Lt => List.drop n l1 ++ (x :: nil) ++ l2
    | Eq => (x :: nil) ++ l2
    | Gt => List.drop (n - Z.of_nat (length l1) - 1)%Z l2
  end.
Proof.
  intros. revert n.
  induction l1; intros n.
  * simpl.
    rewrite Z.sub_0_r.
    destruct (Z.leb_spec n 0), (Z.compare_spec n 0)%Z; try reflexivity; lia.
  * cbn -[Z.add Z.of_nat].
    rewrite IHl1. clear IHl1.
    rewrite Nat2Z.inj_succ, <- Z.add_1_l.
    replace (n - (1 + Z.of_nat (length l1)) - 1)%Z with (n - 2 - Z.of_nat (length l1))%Z by lia.
    replace (n - 1 - (Z.of_nat (length l1)) - 1)%Z with (n - 2 - Z.of_nat (length l1))%Z by lia.
    destruct (Z.leb_spec n 0),
             (Z.compare_spec (n - 1) (Z.of_nat (length l1)))%Z,
             (Z.compare_spec n (1 + Z.of_nat (length l1)))%Z; try reflexivity; lia.
Qed.

Lemma dropGo_spec :
  forall n map lb ub,
  Bounded map lb ub ->
  forall (P : Map e a -> Prop),
  (forall m',
    Bounded m' lb ub /\
    toList m' = List.drop n (toList map) ->
    P m') ->
  P (dropGo n map).
Proof.
  intros ???? HD. revert n.
  induction HD; intros n.
  * intros X HX; apply HX. split.
    + simpl. destruct_match; solve_Bounded e.
    + simpl. do 2 destruct_match; reflexivity.
  * simpl.
    intros X HX; apply HX.
    rewrite toList_Bin.
    unfold op_zlze__ , Ord_Integer___, op_zlze____.
    unfold compare , Ord_Integer___, compare__.
    rewrite size_size.
    apply IHHD1. intros s1' [HBs1' Hs1']. clear IHHD1.
    apply IHHD2. intros s2' [HBs2' Hs2']. clear IHHD2.
    destruct (Z.leb_spec n 0).
    + rewrite drop_neg by assumption. split.
      - solve_Bounded e.
      - rewrite toList_Bin.
        reflexivity.
    + simpl app.
      rewrite drop_app_cons.
      erewrite <- size_spec by eassumption.
      destruct (Z.compare_spec n (size s1)).
      - split.
        ** applyDesc e (@insertMin_Desc e a).
        ** erewrite toList_insertMin  by solve_Precondition e.
           reflexivity.
      - split.
        ** applyDesc e (@link_Desc e a).
        ** erewrite toList_link by solve_Precondition e.
           rewrite Hs1'.
           reflexivity.
      - split.
        ** eapply Bounded_relax_lb; eassumption.
        ** assumption.
Qed.

Lemma toList_drop:
  forall n map lb ub,
  Bounded map lb ub ->
  forall (P : Map e a -> Prop),
  (forall m',
    Bounded m' lb ub /\
    toList m' = List.drop n (toList map) ->
    P m') ->
  P (drop n map).
Proof.
  intros. apply H0.
  unfold drop. fold dropGo.
  unfold op_zgze__ , Ord_Integer___, op_zgze____.
  rewrite size_size.
  destruct (Z.leb_spec (size map) n).
  * split; [solve_Precondition e|].
    rewrite drop_all. reflexivity.
    erewrite <- size_spec by eassumption.
    assumption.
  * eapply dropGo_spec; [solve_Precondition e..|].
    intros. assumption.
Qed.

(** ** Verification of [splitAt] *)

Definition splitAtGo : Int -> Map e a -> (Map e a * Map e a).
Proof.
  let rhs := eval unfold splitAt in (@splitAt e a) in
  match rhs with fun n s => if _ then _ else Datatypes.id (?go _ _)  => exact go end.
Defined.

Lemma splitAtGo_spec :
  forall n s, splitAtGo n s = (takeGo n s, dropGo n s).
Proof.
  intros ??.
  revert n.
  induction s; intros n.
  * simpl.
    repeat destruct_match; try congruence.
  * simpl. repeat destruct_match; reflexivity.
Qed.

End WF.


(*|

========================
Proving safety in TLA
========================

Safety is actually extremely simple in TLA: there's basically one rule, and this file has a few convenient variants.

|*)

From TLA Require Import defs automation.
From TLA Require Import modalities.

Section TLA.

Context [Σ: Type].

Notation exec := (exec Σ).
Notation action := (action Σ).
Notation predicate := (predicate Σ).

Implicit Types (e: exec) (p q: predicate) (a: action).

(* the induction principle from the TLA paper *)
Theorem later_induction (n inv: predicate) :
  (inv ∧ n ⊢ later inv) →
  (inv ∧ □n ⊢ □inv).
Proof.
  unseal.
  destruct H0 as [Hinit Hn].
  induction k.
  - rewrite drop_0 //.
  - change (S k) with (1 + k).
    rewrite -drop_drop.
    apply H; eauto.
Qed.

(* This is a more general induction principle _internal_ to the logic. It's
different from `later_induction` because it requires the implication only for
the "current" execution. *)
Theorem later_induction_internal (n inv: predicate) :
  ⊢ □(inv ∧ n → later inv) → (inv ∧ □n → □inv).
Proof.
  unseal.
  destruct H0 as [Hinit Hn].
  induction k.
  - rewrite drop_0 //.
  - change (S k) with (1 + k).
    apply H; eauto.
Qed.

Theorem tla_init_safety (inv init next safe : predicate) :
  (init ⊢ inv) →
  (inv ∧ next ⊢ later inv) →
  (inv ⊢ safe) →
  ⊢ init ∧ □ next → □ safe.
Proof.
  intros Hinit Hlater Hsafe.
  rewrite Hinit -Hsafe.
  by apply later_induction.
Qed.

Theorem init_safety (inv init: Σ → Prop) (next: action) (safe : Σ → Prop) :
  (∀ s, init s → inv s) →
  (∀ s s', inv s → next s s' → inv s') →
  (∀ s, inv s → safe s) →
  ⊢ state_pred init ∧ □ ⟨next⟩ → □ (state_pred safe).
Proof.
  intros Hinit Hinv Hsafe.
  apply (tla_init_safety (state_pred inv)); unseal.
Qed.

Theorem init_invariant (init: Σ → Prop) (next: action) (inv: Σ → Prop) :
  (∀ s, init s → inv s) →
  (∀ s s', inv s → next s s' → inv s') →
  state_pred init ∧ □ ⟨next⟩ ⊢ □ state_pred inv.
Proof.
  intros Hinit Hinv.
  apply (init_safety inv); auto.
Qed.

End TLA.


From Cofq.SystemF Require Import
     SystemFDefinitions
     SystemFUtils
     SystemFShow. (* Need show for errors in evaluation *)

From Cofq.Utils Require Import Utils.
From Cofq.Show Require Import ShowUtils.

From Coq Require Import
     Lia
     String.

From ExtLib Require Import
     Structures.Monads
     Structures.Functor
     Eqv.

Import MonadNotation.
Local Open Scope monad_scope.

From ITree Require Import
     Basics
     ITree
     Interp.Recursion
     Events.Exception.

From Vellvm.Utils Require Import Util.

Section Substitution.
  (* Lift by 2 because fixpoint has a argument in addition to referring to itself *)
  Fixpoint term_lift {I} `{FInt I} (n : N) (term : Term) : Term :=
    match term with
    | Var x =>
      if N.ltb x n
      then Var x
      else Var (x + 2)
    | Ann term' type => Ann (term_lift n term') type
    | Fix fix_type arg_type fbody => Fix fix_type arg_type (term_lift (n+2) fbody)
    | App e1 e2 => App (term_lift n e1) (term_lift n e2)
    | TAbs e => TAbs (term_lift n e)
    | TApp e t => TApp (term_lift n e) t
    | Tuple es => Tuple (map (term_lift n) es)
    | ProjN i es => ProjN i (term_lift n es)
    | Num x => Num x
    | If0 c e1 e2 => If0 (term_lift n c) (term_lift n e1) (term_lift n e2)
    | Op op e1 e2 => Op op (term_lift n e1) (term_lift n e2)
    end.

  Fixpoint type_lift (n : N) (τ : FType) : FType :=
    match τ with
    | Arrow τ1 τ2 => Arrow (type_lift n τ1) (type_lift n τ2)
    | Prod τs => Prod (map (type_lift n) τs)
    | TForall τ' => TForall (type_lift (N.succ n) τ')
    | TVar x => if N.ltb x n then TVar x else TVar (x + 1)
    | IntType => τ
    end.

  Fixpoint term_subst {I} `{FInt I} (v : VarInd) (body arg : Term) : Term :=
    match body with
    | Var x =>
      if N.eqb x v
      then arg
      else Var x
    | Fix fix_type arg_type fbody =>
      Fix fix_type arg_type (term_subst (v + 2) fbody (term_lift 0 arg))
    | Ann e t => Ann (term_subst v e arg) t
    | App e1 e2 => App (term_subst v e1 arg) (term_subst v e2 arg)
    | TAbs e => TAbs (term_subst v e arg)
    | TApp e t => TApp (term_subst v e arg) t
    | Tuple es => Tuple (map (fun e => term_subst v e arg) es)
    | ProjN i e => ProjN i (term_subst v e arg)
    | Num x => Num x
    | If0 c e1 e2 => If0 (term_subst v c arg) (term_subst v e1 arg) (term_subst v e2 arg)
    | Op op e1 e2 => Op op (term_subst v e1 arg) (term_subst v e2 arg)
    end.

  Lemma type_lift_type_size:
    forall τ n,
      type_size (type_lift n τ) = type_size τ.
  Proof.
    induction τ; intros sz; cbn;
      repeat match goal with
               H: _ |- _ =>
               rewrite H
             end; eauto.

    destruct (x <? sz)%N; auto.

    rewrite map_map.
    erewrite map_ext_in; eauto.
  Qed.

  Obligation Tactic := try Tactics.program_simpl; try solve [cbn; try lia | repeat split; try solve [intros; discriminate]].
  Program Fixpoint type_subst_in_type (v : TypeInd) (τ : FType) (arg : FType) {measure (type_size τ)} : FType :=
    match τ with
    | TVar x =>
      if N.eqb x v
      then arg
      else if N.ltb v x
           then TVar (x-1)
           else TVar x
    | Arrow τ1 τ2 => Arrow (type_subst_in_type v τ1 arg) (type_subst_in_type v τ2 arg)
    | Prod τs =>
      Prod (map_In τs (fun τ HIn => type_subst_in_type v τ arg))
    | TForall τ' => TForall (type_subst_in_type (v+1) τ' (type_lift 0 arg))
    | IntType => IntType
    end.
  Next Obligation.
    cbn.
    pose proof (list_sum_map type_size τ τs HIn).
    lia.
  Qed.

  Fixpoint type_subst {I} `{FInt I} (v : TypeInd) (e : Term) (arg_type : FType) : Term
    := match e with
       | TAbs e => TAbs (type_subst (v+1) e (type_lift 0 arg_type))
       | TApp e τ => TApp (type_subst v e arg_type) (type_subst_in_type v τ arg_type)
       | Fix fτ τ body => Fix (type_subst_in_type v fτ arg_type) (type_subst_in_type v τ arg_type) (type_subst v body arg_type)
       | App e1 e2 => App (type_subst v e1 arg_type) (type_subst v e2 arg_type)
       | If0 c e1 e2 => If0 (type_subst v c arg_type) (type_subst v e1 arg_type) (type_subst v e2 arg_type)
       | Op op e1 e2 => Op op (type_subst v e1 arg_type) (type_subst v e2 arg_type)
       | Tuple es => Tuple (map (fun e => type_subst v e arg_type) es)
       | ProjN i e => ProjN i (type_subst v e arg_type)
       | Ann e τ => Ann (type_subst v e arg_type) (type_subst_in_type v τ arg_type)
       | Num x => Num x
       | Var x => Var x
       end.

  Definition app_fix {I} `{FInt I} (fix_type arg_type : FType) (body : Term) (arg : Term) : Term :=
    term_subst 1 (term_subst 0 body (Fix fix_type arg_type body)) arg.
End Substitution.

(** Single-step semantics for SystemF, useful for testing. *)
Section SingleStep.
  Obligation Tactic := try Tactics.program_simpl; try solve [cbn; try lia | repeat split; try solve [intros; discriminate]].
  Program Fixpoint step {I} `{FInt I} (e : Term) {measure (term_size e)} : (unit + Term) :=
    match e with
    | Ann e t => step e
    | Fix fix_type arg_type body => inl tt
    | App (Fix fix_type arg_type body) arg =>
      inr (app_fix fix_type arg_type body arg)
    | App e1 e2 =>
      e1v <- step e1;;
      inr (App e1v e2)
    | Op op (Num xn) (Num yn) =>
      inr (Num (eval_op op xn yn))
    | Op op (Num xn) y =>
      yv <- step y ;;
      inr (Op op (Num xn) yv)
    | Op op x y =>
      xv <- step x ;;
      inr (Op op xv y)
    | TApp (TAbs e) arg_type =>
      inr (type_subst 0 e arg_type)
    | TApp e t =>
      e' <- step e;;
      inr (TApp e' t)
    | ProjN i (Tuple es) =>
      match nth_error es (N.to_nat i) with
      | Some e => step e
      | None => inl tt
      end
    | ProjN i e =>
      ev <- step e;;
      inr (ProjN i ev)
    | If0 (Num x) e1 e2 =>
      if eq x zero
      then inr e1
      else inr e2
    | If0 c e1 e2 =>
      cv <- step c;;
      inr (If0 cv e1 e2)
    | _ => inl tt
    end.
  Next Obligation.
    cbn.
    symmetry in Heq_anonymous.
    apply nth_error_In in Heq_anonymous.
    pose proof (list_sum_map term_size e es Heq_anonymous).
    lia.
  Qed.

  Definition step' {I} `{FInt I} (v : unit + Term) : (unit + Term)
    := match v with
       | inl tt => inl tt
       | inr t => step t
       end.

  Fixpoint multistep {I} `{FInt I} (n : nat) (v : Term) : Term
    := match n with
       | O => v
       | S n =>
         match step v with
         | inl tt => v
         | inr t => multistep n t
         end
       end.

  (* Returns a term only if it's fully evaluated *)
  Fixpoint multistep' {I} `{FInt I} (n : nat) (v : Term) : option Term
    := match n with
       | O => None
       | S n =>
         match step v with
         | inl tt => Some v
         | inr t => multistep' n t
         end
       end.
End SingleStep.

(** Denotation of SystemF in terms of itrees. *)
Section Denotation.
  Definition Failure := exceptE string.

  Definition eval_body {I} `{FInt I} `{Show I} (e : Term) : itree (callE Term Term +' Failure) Term :=
    match e with
    | Ann e t => call e
    | App e1 e2 =>
      e1v <- call e1;;
      e2v <- call e2;;
      match e1v with
      | Fix fix_type arg_type body =>
        call (app_fix fix_type arg_type body e2v)
      | _ => throw ("ill-typed application: " ++ show e)
      end
    | TApp e t =>
      e' <- call e;;
      match e' with
      | TAbs body =>
        call (type_subst 0 body t)
      | _ => throw ("ill-typed type application: " ++ show (TApp e' t))
      end
    | ProjN i es =>
      es' <- call es;;
      match es' with
      | Tuple xs =>
        match nth_error xs (N.to_nat i) with
        | Some e => call e
        | None => throw ("tuple projection out of bounds: " ++ show e)
        end
      | _ => throw ("ill-typed tuple projection: " ++ show e)
      end
    | If0 c e1 e2 =>
      cv <- call c;;
      match cv with
      | Num x =>
        if eq x zero
        then call e1
        else call e2
      | _ => throw ("ill-typed if0: " ++ show e)
      end
    | Op op e1 e2 =>
      e1v <- call e1;;
      e2v <- call e2;;
      match e1v, e2v with
      | Num x, Num y =>
        ret (Num (eval_op op x y))
      | _, _ => throw ("ill-typed operation: " ++ show e)
      end
    | Num x => ret e
    | TAbs x => ret e
    | Tuple xs =>
      xs' <- map_monad call xs;;
      ret (Tuple xs')
    | Fix fix_type arg_type body => ret e
    | Var x => ret e
    end.

  Definition eval {I} `{FInt I} `{Show I} : Term -> itree Failure Term :=
    rec eval_body.
End Denotation.


Generalizable All Variables.
Set Implicit Arguments.
Unset Strict Implicit.
Unset Printing Implicit Defensive.
Set Reversible Pattern Implicit.
Set Primitive Projections.
Set Universe Polymorphism.

Require Export setoid.
Require Import ProofIrrelevance.
Require Import ChoiceFacts.

Declare Scope cat_scope.
Open Scope setoid_scope.
Open Scope cat_scope.

Obligation Tactic :=
  (try now intros x); (try now split; intros x); intros; (try apply \ISE).

Inductive Composables {obj hom : Setoid} (dom cod : Map hom obj) :=
  ligature : forall f g, cod f == dom g -> Composables _ _.
Notation "$[ g , f | H ]" := (@ligature _ _ _ _ f g H)
  (at level 0, format "$[ g ,  f  |  H ]") : cat_scope.
Notation "{ 'o' dom , cod }" := (@Composables _ _ dom cod)
  (at level 0, dom, cod at level 99,
  format "{ 'o'  dom ,  cod }") : cat_scope.

Definition composables_eq {obj hom : Setoid} {dom cod : Map hom obj}
  (c c0 : Composables dom cod) :=
  let (f, g, H) := c in let (f0, g0, H0) := c0 in f == f0 /\ g == g0.
#[global] Hint Unfold composables_eq : eq.

Program Canonical Structure ComposablesSetoid
  {obj hom : Setoid} (dom cod : Map hom obj) :=
  [ _ | ==: @composables_eq _ _ dom cod ].
Next Obligation.
  split.
  - intros [f g H]. now split.
  - intros [f g H] [f0 g0 H0] [E E0]. now split.
  - intros [f g H] [f0 g0 H0] [f1 g1 H1] [E E0] [E1 E2].
    split; now (rewrite E || rewrite E0).
Defined.
Notation "[ 'o' dom , cod ]" := (@ComposablesSetoid _ _ dom cod)
  (at level 0, dom, cod at level 99,
  format "[ 'o'  dom ,  cod ]") : cat_scope.

Class IsCategory {obj hom : Setoid} {dom cod : Map hom obj}
  (comp : Map [o dom, cod] hom) (id : obj -> hom) :=
{
  comp_dom : forall f g H, dom (comp $[g, f | H]) == dom f;
  comp_cod : forall f g H, cod (comp $[g, f | H]) == cod g;
  comp_assoc : forall f g h H H0 H1 H2,
    comp $[h, comp $[g, f | H] | H0] ==
    comp $[comp $[h, g | H1], f | H2];
  id_dom : forall A, dom (id A) == A;
  id_cod : forall A, cod (id A) == A;
  comp_idr : forall f H, comp $[f, id (dom f) | H] == f;
  comp_idl : forall f H, comp $[id (cod f), f | H] == f;
}.

Structure Category := {
  catobj :> Setoid;
  cathom : Setoid;
  catdom : Map cathom catobj;
  catcod : Map cathom catobj;
  homcomp : Map [o catdom, catcod] cathom;
  catid : Map catobj cathom;

  catprf :> IsCategory homcomp catid
}.
#[global] Existing Instance catprf.

Notation "[ 'dom:' dom , 'cod:' cod , 'o:' comp , '1:' id ]"
  := (@Build_Category _ _ dom cod comp id _)
  (at level 0, dom, cod, comp, id at level 99) : cat_scope.
Notation "[ 'dom' F => A , 'cod' G => B , 'comp' C => P , 'id' D => Q ]"
  := [ dom: map F => A, cod: map G => B, o: map C => P, 1: map D => Q ]
  (at level 0, F binder, G binder, C binder, D binder,
  A, B, P, Q at level 99) : cat_scope.
Notation "'dom' f" := (@catdom _ f)
  (at level 60, right associativity) : cat_scope.
Notation "'cod' f" := (@catcod _ f)
  (at level 60, right associativity) : cat_scope.
Notation "g o[ 'by' H ] f" := (@homcomp _ $[g, f | H])
  (at level 60, H at level 99, right associativity,
  format "g  o[ 'by'  H ]  f") : cat_scope.
Notation "g o[] f" := (g o[by _] f)
  (at level 60, right associativity, format "g  o[]  f") : cat_scope.
Notation "1_ A" := (@catid _ A)
  (at level 0, right associativity, format "1_ A") : cat_scope.

Section uouo.
Variables obj hom : Setoid.
Variables Dom Cod : Map hom obj.
Check [o Dom, Cod].
Variables X : Category.
Variables f g : cathom X.
Variables H : cod f == dom g.
Check (g o[by H] f).
End uouo.

Program Definition homDC {X : Category} :=
  dmap A B => [ f | dom f == A /\ cod f == B ].
Next Obligation.
  intros f f0 E. split; intros [E0 E1]; split; now rewrite2 E.
Defined.
Next Obligation.
  intros A A0 E B B0 E0 f f0 E1. split; intros [H H0]; split;
  rewrite2 E1; try (now rewrite H); now rewrite H0.
Defined.
Notation "A ~> B" := (@homDC _ A B)
  (at level 90, no associativity) : cat_scope.

Program Definition homDC_comp {X : Category} {A B C : X}
  : Dymap (A ~> B) (B ~> C) (A ~> C)
  := dmap f g => $[g o[by _] f, _].
Next Obligation.
  destruct f as [f [H H0]]. destruct g as [g [H1 H2]].
  simpl. now rewrite H0.
Defined.
Next Obligation.
  destruct f as [f [H H0]]. destruct g as [g [H1 H2]].
  split; simpl in *; now (rewrite comp_dom || rewrite comp_cod).
Defined.
Next Obligation.
  intros f f0 E g g0 E0. simpeq_all. apply map_equal. split; now simpl.
Defined.
Notation "g 'o' f" := (@homDC_comp _ _ _ _ f g)
  (at level 60, right associativity) : cat_scope.

Program Definition idDC {X : Category} (A : X) : A ~> A
  := $[1_A, _].
Next Obligation. split; now (rewrite id_dom || rewrite id_cod). Defined.
Notation "1_[ A ]" := (@idDC _ A)
  (at level 0, format "1_[ A ]") : cat_scope.

Program Definition hom_homDC `(f : cathom X) : dom f ~> cod f
  := $[f, _].

Section CatTheory.
  Context {X : Category}.
  Implicit Types (A B C D : X) (f g : cathom X).

  Lemma compeq f g H H0 : g o[by H] f == g o[by H0] f.
  Proof. apply map_equal. now split. Qed.

  Lemma comp1fH f H : 1_(cod f) o[by H] f == f.
  Proof. apply comp_idl. Qed.

  Lemma comp1F A B (F : A ~> B) : 1_[B] o F == F.
  Proof.
    case F as [f [H H0]]. simpeq.
    assert (cod f == dom 1_(cod f)) by now rewrite id_dom.
    rewrite <-(comp_idl H1) at 3. apply map_equal.
    split; intuition. now apply map_equal.
  Qed.

  Lemma compF1 A B (F : A ~> B) : F o 1_[A] == F.
  Proof.
    case F as [f [H H0]]. simpeq.
    assert (cod 1_(dom f) == dom f) by now rewrite id_cod.
    rewrite <-(comp_idr H1) at 6. apply map_equal.
    split; intuition. now apply map_equal.
  Qed.

  Lemma compA A B C D (F : A ~> B) (G : B ~> C) (H : C ~> D)
    : H o G o F == (H o G) o F.
  Proof. apply comp_assoc. Qed.
End CatTheory.

Class IsIsopair {X : Category} {A B : X} (f : A ~> B) (g : B ~> A) := {
  isodom : g o f == 1_[A];
  isocod : f o g == 1_[B]
}.

Definition IsIsomorphic (X : Category) (A B : X) :=
  exists (f : A ~> B) (g : B ~> A), IsIsopair f g.

Program Definition IsomorphSetoid (X : Category) :=
  [ X | ==: @IsIsomorphic X ].
Next Obligation.
  split; intros A.
  - exists 1_[A], 1_[A]. split; now rewrite comp1F.
  - intros B [f [g [H H0]]]. exists g, f. now split.
  - intros B C [f [g [H H0]]] [f0 [g0 [H1 H2]]].
    exists (f0 o f), (g o g0). split.
  + now rewrite <-compA, (compA f), H1, comp1F.
  + now rewrite <-compA, (compA g0), H0, comp1F.
Defined.
Notation "A === B" := (A == B in @IsomorphSetoid _)
  (at level 70, no associativity) : cat_scope.

Program Definition OppCat (X : Category) :=
  [ dom: catcod X, cod: catdom X,
    o: map c => let (f, g, H) := c in f o[by _] g,
    1: catid X ].
Next Obligation.
  intros [f g H] [f0 g0 H0] [E E0]. apply map_equal. now split.
Defined.
Next Obligation.
  split; intuition; simpl.
  - apply comp_cod.
  - apply comp_dom.
  - symmetry. apply comp_assoc.
  - apply id_cod.
  - apply id_dom.
  - apply comp_idl.
  - apply comp_idr.
Defined.
Notation "C ^op" := (@OppCat C)
  (at level 40, left associativity, format "C ^op") : cat_scope.

Class IsFunctor (X Y : Category) (fobj : X -> Y)
  (fhom : Map (cathom X) (cathom Y)) :=
{
  compF_dom : forall f, dom (fhom f) == fobj (dom f);
  compF_cod : forall f, cod (fhom f) == fobj (cod f);
  compF_morph : forall f g H H0,
    fhom (g o[by H] f) == (fhom g) o[by H0] (fhom f);
  idF_morph : forall A, fhom 1_A == 1_(fobj A)
}.

Structure Functor (X Y : Category) := {
  funobj :> Map X Y;
  funhom : Map (cathom X) (cathom Y);
  funprf :> IsFunctor funobj funhom
}.
#[global] Existing Instance funprf.

Notation "X --> Y" := (@Functor X Y) : cat_scope.
Notation "F ':o' f" := (@funhom _ _ F f)
  (at level 60, right associativity) : cat_scope.
Notation "[ 'obj:' F , 'hom:' Ff ]" :=
  (@Build_Functor _ _ F Ff _)
  (at level 0, F, Ff at level 99, no associativity) : cat_scope.
Notation "[ 'obj' A => FA , 'hom' f => Ff ]" :=
  [ obj: map A => FA , hom: map f => Ff ]
  (at level 0, A binder, f binder,
  FA, Ff at level 99, no associativity)
  : cat_scope.

Program Canonical Structure FunctorSetoid X Y :=
  [ ==: (F : X --> Y) G => forall f, F :o f == G :o f ].
Next Obligation.
  split; try now intros F.
  intros F G H H0 H1 f. now rewrite H0, H1.
Defined.
Notation "X [-->] Y" := (@FunctorSetoid X Y)
  (at level 55) : cat_scope.

Program Definition funhomDC `(F : X --> Y) (A B : X)
  : Map (A ~> B) (F A ~> F B) := map f => $[F :o f, _].
Next Obligation.
  destruct F as [Fo Ff H]. destruct f as [f [H0 H1]].
  split; simpl; [rewrite compF_dom | rewrite compF_cod];
  now apply map_equal.
Defined.
Next Obligation.
  intros f f0 E. simpeq_all. now apply map_equal.
Defined.
Notation "F :[o] f" := (@funhomDC _ _ F _ _ f)
  (at level 60, right associativity) : cat_scope.

Section FunctorTheory.
  Context {X Y : Category} {F : X --> Y}.
  Implicit Types (A B C : X).

  Lemma compFM A B C (f : A ~> B) (g : B ~> C) :
    F :[o] (g o f) == (F :[o] g) o F :[o] f.
  Proof. apply compF_morph. Qed.

  Lemma idFM A : F :[o] 1_[A] == 1_[F A].
  Proof. apply idF_morph. Qed.
End FunctorTheory.

Definition eqSetoid X := [ X | ==: @eq X ].
Notation "[ 'eqS' X ]" := (@eqSetoid X)
  (at level 0, format "[ 'eqS'  X ]") : cat_scope.

Structure Carto {X} (F : X -> X -> Setoid) := {
  cdom : [eqS X];
  ccod : [eqS X];
  cfun :> F cdom ccod
}.

Inductive cfun_eq `{F : X -> X -> Setoid}
  A B (f : F A B) : forall C D, F C D -> Prop :=
  cfun_eqref : forall g : F A B, f == g -> cfun_eq f g.

Program Canonical Structure CartoSetoid `(F : X -> X -> Setoid) :=
  [ ==: (f : Carto F) G => @cfun_eq _ _ (cdom f) (ccod f) (cfun f) _ _ G ].
Next Obligation.
  split; try now intros [f]; try now intros [g] [H].
  intros [f] [g] [h] [H0 H1] [H2 H3]. split. simpl in *.
  now rewrite H1, H3.
Defined.

(* Type Category *)

Program Definition funcSetoid X Y :=
  [ X -> Y | ==: fun f g => forall x, f x = g x ].
Next Obligation.
  split; intuition. intros f g h E E0 x. now rewrite E.
Defined.

Definition funcS_comp `(f : funcSetoid X Y) `(g : funcSetoid Z W)
  : Y = Z -> funcSetoid X W.
Proof.
  intros H. revert f g. case H. intros f g.
  refine (fun x => g (f x)).
Defined.

Program Definition TypeCat :=
  [ dom: map F => cdom F, cod: map F => ccod F,
    o: map p => let (f,g,H) := p in Build_Carto (funcS_comp (cfun f) (cfun g) H),
    1: map X => (@Build_Carto _ _ X X (fun x => x)) ].
Next Obligation. now intros [d c f] [d0 c0 f0] [E H]. Defined.
Next Obligation. now intros [d c f] [d0 c0 f0] [E H]. Defined.
Next Obligation.
  intros [[df cf f] [dg cg g] H] [[df0 cf0 f0] [dg0 cg0 g0] H0] [E E0].
  simpl in *. destruct E, E0, H.
  rewrite (proof_irrelevance _ H0 (eq_refl _)).
  split. compute. intros x. now rewrite H2, H1.
Defined.
Next Obligation.
  intros X X0 E. destruct E. split. now simpl.
Defined.
Next Obligation.
  split; try intros [df cf f] [dg cg g] H; try now simpl in *.
  2,3: intros [df cf f] H;
  rewrite (proof_irrelevance _ H (eq_refl _));
  split; now compute.
  destruct H as [dh ch h]. intros H H0 H1 H2. simpl in *.
  destruct H, H0.
  rewrite (proof_irrelevance _ H1 (eq_refl _)).
  rewrite (proof_irrelevance _ H2 (eq_refl _)).
  split. now compute.
Defined.

(* CAT *)

Program Definition Functor_comp {X Y Z : Category}
  : Dymap (X --> Y) (Y --> Z) (X --> Z) := dmap F G =>
  [ obj (A : X) => G (F A), hom (f : cathom X) => G :o F :o f ].
Next Obligation. intros A A0 E. now rewrite E. Defined.
Next Obligation. intros f f0 E. now rewrite E. Defined.
Next Obligation.
  split; simpl in *; intros f.
  - now rewrite !compF_dom.
  - now rewrite !compF_cod.
  - intros g H H0. now rewrite !compF_morph.
  - now rewrite !idF_morph. Unshelve.
    now rewrite compF_cod, compF_dom, H.
Defined.
Next Obligation.
  intros F F0 E G G0 E0 f. simpl in *. now rewrite E0, E.
Defined.

Definition FunctorC_comp `(f : X [-->] Y) `(g : Z [-->] W)
  : Y = Z -> X [-->] W.
Proof.
  intros H. revert f g. case H. intros f g.
  refine (Functor_comp f g).
Defined.

Program Definition Functor_id (X : Category)
  := [ obj A => A, hom f => f ].
Next Obligation.
  split; intros f; simpl; intuition. now apply map_equal.
Defined.

Program Canonical Structure CAT :=
  [ dom F => cdom F, cod F => ccod F,
    comp p => let (f,g,H) := p in Build_Carto (FunctorC_comp (cfun f) (cfun g) H),
    id X => (@Build_Carto _ _ X X (Functor_id X)) ].
Next Obligation.
  intros [dF cF F] [dF0 cF0 F0] [E]. now simpl in *.
Defined.
Next Obligation.
  intros [dF cF F] [dF0 cF0 F0] [E]. now simpl in *.
Defined.
Next Obligation.
  intros [[df cf f] [dg cg g] H] [[df0 cf0 f0] [dg0 cg0 g0] H0] [E E0].
  simpl in *. destruct E, E0, H. split. simpl in *. intros f0.
  rewrite (proof_irrelevance _ H0 eq_refl). simpl.
  now rewrite H1, H2.
Defined.
Next Obligation. intros X Y E. now destruct E. Defined.
Next Obligation.
  split; try intros [df cf f] [dg cg g] H; try now simpl in *.
  2,3: intros [df cf f] H;
  rewrite (proof_irrelevance _ H (eq_refl _));
  split; now compute.
  destruct H as [dh ch h]. intros H H0 H1 H2. simpl in *.
  destruct H, H0.
  rewrite (proof_irrelevance _ H1 (eq_refl _)).
  rewrite (proof_irrelevance _ H2 (eq_refl _)).
  split. now compute.
Defined.

Class IsFaithful `(F : X --> Y) := {
  faith : forall A B (f g : A ~> B), F :o f == F :o g -> f == g
}.

Class IsFull `(F : X --> Y) := {
  full : forall A B (f : F A ~> F B),
    exists (g : A ~> B), F :o g == f
}.

Class IsNattrans {X Y} (F G : X --> Y) (c : Map X (cathom Y)) := {
  nat_dom : forall A, dom (c A) == F A;
  nat_cod : forall A, cod (c A) == G A;
  natural : forall f H H0,
    (c (cod f)) o[by H] (F :o f) == (G :o f) o[by H0] (c (dom f))
}.

Structure Nattrans {X Y} (F G : X --> Y) := {
  component :> Map X (cathom Y);
  ntprf :> IsNattrans F G component
}.
#[global] Existing Instance ntprf.

Notation "F ==> G" := (@Nattrans _ _ F G) : cat_scope.
Notation "[ ==>: C ]" := (@Build_Nattrans _ _ _ _ C _)
  (at level 0, C at level 99) : cat_scope.
Notation "[ 'nt' A => P ]" := [ ==>: fun A => P ]
  (at level 0, A binder, P at level 99) : cat_scope.

Lemma ntcomm X Y (F G : X --> Y) (a : F ==> G) A B (f : A ~> B)
  : (a B) o (F :[o] f) == (G :[o] f) o (a A).
Proof. apply ntprf. Qed.

Definition nt_eq {X Y} (F G : X --> Y) (a b : F ==> G) :=
  forall A, a A == b A.
#[global] Hint Unfold nt_eq : eq.
Program Canonical Structure NattransSetoid {X Y}
  (F G : X --> Y) := [ _ | ==: @nt_eq _ _ F G ].
Next Obligation.
  split; try now intros a. intros a b c H H0 A.
  simpeq_all. now rewrite H, H0.
Defined.
Notation "F [==>] G" := (@NattransSetoid _ _ F G)
  (at level 55, right associativity) : cat_scope.

Program Definition Nattrans_vcomp {X Y} {F G H : X --> Y}
  : Dymap (F ==> G) (G ==> H) (F ==> H) := dmap a b =>
  [ nt A => (b A) o (a A) ].
Next Obligation.
  split; intros. now rewrite <-compA, ntcomm, compA, ntcomm, compA.
Defined.
Next Obligation.
  intros a a0 E b b0 E0 A. simpeq_all.
  apply map_equal. split; simpl; now (rewrite E || rewrite E0).
Defined.
Notation "b |o| a" := (@Nattrans_vcomp _ _ _ _ _ a b)
  (at level 53, right associativity) : cat_scope.

Definition NatrransC_vcomp {X Y} {F F0 F1 F2 : X --> Y}
  (a : F ==> F0) (b : F1 ==> F2) : F0 == F1 -> F0 ==> F2.
Proof.
  intros H. rewrite H.

`(f : X [-->] Y) `(g : Z [-->] W)
  : Y = Z -> X [-->] W.
Proof.
  intros H. revert f g. case H. intros f g.
  refine (Functor_comp f g).
Defined.

Program Definition Functor_id (X : Category)
  := [ obj A => A, hom f => f ].
Next Obligation.
  split; intros f; simpl; intuition. now apply map_equal.
Defined.

Program Canonical Structure FunctorCat (X Y : Category) :=
  [ dom F => cdom F, cod F => ccod F,
    comp p => let (f, g, H) := p in Build_Carto (cfun f |o| cfun g),
    id X => (@Build_Carto _ _ X ) ]
[ dom F => cdom F, cod F => ccod F,
comp p => let (f,g,H) := p in Build_Carto (FunctorC_comp (cfun f) (cfun g) H),
id X => (@Build_Carto _ _ X X (Functor_id X)) ].
  [ dom  ]
  [ hom (F : X --> Y) G => Nattrans F G, comp F G H a b => b |o| a,
    id F => [ nt A => 1_(F A) ] ].
Next Obligation. apply Nattrans_vcomp_obligation_2. Defined.
Next Obligation. split; intros. now rewrite compf1, comp1f. Defined.
Next Obligation.
  split; intros; simpeq; intros.
  - now rewrite compA.
  - now rewrite compf1.
  - now rewrite comp1f.
Defined.
Notation "[ X , Y ]" := (@FunctorCat X Y)
  (at level 0, X, Y at level 99) : cat_scope.

Lemma natiso_cmpiso X Y (F G : X --> Y) (a : F ==> G) :
  DependentFunctionalChoice ->
  IsIsomorphism a == forall A, IsIsomorphism (a A).
Proof.
  intros DFC. unfold IsIsomorphism. split; simpeq_all.
  - intros [g [E E0]] A. exists (g A). split;
    (apply E || apply E0).
  - intros H.
    pose (DFC X (fun A => G A ~> F A) (fun A g => g o a A == 1_(F A) /\ a A o g == 1_(G A))).
    simpl in *. destruct (e H) as [g H0].
    assert (IsNattrans g).
    { split. intros. now rewrite <-(compf1 (G :o f)),
      <-(proj2 (H0 A)), (compA (g A)), <-natural, !compA,
      (proj1 (H0 B)), comp1f. }
    exists (Build_Nattrans H1). split; intros A; simpl in *;
    apply (proj1 (H0 A)) || apply (proj2 (H0 A)).
Qed.

Class IsEquivCat `(F : X --> Y) (G : Y --> X) (eta : 1_X ==> G :o: F)
  (xi : F :o: G ==> 1_Y) :=
{
  ecetaiso : IsIsomorphism eta;
  ecxiiso : IsIsomorphism xi
}.

Structure EquivCat (X Y : Category) := {
  ecfwd : X --> Y;
  ecrev : Y --> X;
  eceta : 1_X ==> ecrev :o: ecfwd;
  ecxi : ecfwd :o: ecrev ==> 1_Y;
  ecprf :> IsEquivCat eceta ecxi
}.
#[global] Existing Instance ecprf.

Program Definition Nattrans_hcomp {X Y Z} {F G : X --> Y} {F0 G0 : Y --> Z}
  : Dymap (F ==> G) (F0 ==> G0) (F0 :o: F ==> G0 :o: G) := dmap a b =>
  [ nt A => b (G A) o (F0 :o a A) ].
Next Obligation.
  split. intros A B f. simpl. rewrite !ntcomm.
  rewrite <-!compA. rewrite !ntcomm. rewrite compA.
  now rewrite <-compFC, !ntcomm, compFC, compA.
Defined.
Next Obligation.
  intros a a0 E b b0 E0 A. simpeq_all. now rewrite E, E0.
Defined.
Notation "b =o= a" := (@Nattrans_hcomp _ _ _ _ _ _ _ a b)
  (at level 53, right associativity) : cat_scope.

Program Definition nt1_1o1 {X : Category} : 1_X ==> 1_X :o: 1_X :=
  [ nt A => 1_A ].
Next Obligation.
  split. intros A B f. simpl. now rewrite comp1f, compf1.
Defined.

Program Definition nt1o1_1 {X : Category} : 1_X :o: 1_X ==> 1_X :=
  [ nt A => 1_A ].
Next Obligation.
  split. intros A B f. simpl. now rewrite comp1f, compf1.
Defined.

(* Program Definition nt_refconv {X Y} {F F0 G G0 : X --> Y}
  (E : F == F0) (E0 : G == G0) (a : F ==> G) : F0 ==> G :=
  [ nt A => a A ]. *)

(* Program Definition nt_reflconv0 {X Y X} {F : X --> Y} {G : Y --> X}
  {F0 : Y --> Z} {G0 : Z --> Y}
  (eta0 : 1_Y ==> G0 :o: F0) *)

Program Definition nt_revlconv {X Y} {F F0 : X --> Y} (H : F == F0)
  : F ==> F0.
rewrite H.

Program Definition EquivCatSetoid :=
  [ ==: (X : Category) Y => exists (F : X --> Y) (G : Y --> X)
    (eta : 1_X ==> G :o: F) (xi : F :o: G ==> 1_Y),
    IsEquivCat eta xi ].
Next Obligation.
  split.
  - intros X. exists 1_X, 1_X, nt1_1o1, nt1o1_1.
    split; unfold IsIsomorphism; [exists nt1o1_1 | exists nt1_1o1];
    split; simpeq; intros A; now rewrite comp1f.
  - intros X Y [F [G [eta [xi [[g [H H0]] [g0 [H1 H2]]]]]]].
    exists G, F, g0, g. split; [exists xi | exists eta]; split; intuition.
  - intros X Y Z [F [G [eta [xi H]]]] [F0 [G0 [eta0 [xi0 H0]]]].
    exists (F0 :o: F), (G :o: G0).
    pose (1_G =o= eta0 =o= 1_F). pose (s o eta). exists s. simpl in s. pose (eta0 o eta).


From mathcomp Require Import all_ssreflect all_algebra.
Require Import lens.

Set Implicit Arguments.
Unset Strict Implicit.
Unset Printing Implicit Defensive.

Import GRing.Theory.
Local Open Scope ring_scope.

Reserved Notation "f \v g" (at level 50, format "f  \v  g").
Reserved Notation "f =e g" (at level 70).
Reserved Notation "M1 '*t' M2" (at level 50).

(* Reduce a linear form *)
Definition linE :=
  (mulr0,mul0r,mulr1,mul1r,addr0,add0r,subr0,oppr0,scale0r,scale1r).

Section tensor_space.
Variables (I : finType) (dI : I) (R : comRingType).
Local Notation merge_indices := (merge_indices dI).

Definition tpower n T := {ffun n.-tuple I -> T}.
Definition morfun m n := forall T : lmodType R, tpower m T -> tpower n T.
Definition mor m n :=
  forall T : lmodType R, {linear tpower m T -> tpower n T}.
Definition tmatrix m n := tpower m (tpower n R^o).
Notation tsquare n := (tmatrix n n).
Notation endo n := (mor n n).
Notation endofun n := (morfun n n).

Definition tpcast n m T (H : n = m) (v : tpower n T) : tpower m T :=
  [ffun vi => v (cast_tuple vi (esym H))].

Lemma tpcastE T n v (H : n = n) : tpcast (T:=T) H v = v.
Proof. by apply/ffunP => vi; rewrite !ffunE cast_tupleE. Qed.

Lemma tpcastK T n m (H : n = m) (t : tpower n T) :
  tpcast (esym H) (tpcast H t) = t.
Proof. by apply/ffunP => v; rewrite !ffunE; f_equal; apply/val_inj. Qed.

Lemma tpcast_linear (T : lmodType R) n m (H : n = m) : linear (tpcast (T:=T) H).
Proof. move=> x y z /=; apply/ffunP => vi; by rewrite !ffunE. Qed.

Definition mor_tpcast n m (H : n = m) : mor n m :=
  fun T : lmodType R => Linear (tpcast_linear (T:=T) H).
Canonical mor_tpcast.

(* Actually, need the property (naturality)
 forall (f : endo m) (T1 T2 : lmodType R) (h : {linear T1 -> T2}),
   map h \o f T1 = f T2 \o map h
which is equivalent to the fact f = nvendo M for a square matrix M : tsquare m.
*)
Definition map_tpower m T1 T2 (f : T1 -> T2) (nv : tpower m T1)
  : tpower m T2 := [ffun v : m.-tuple I => f (nv v)].

Definition naturality m n (f : mor m n) :=
  forall (T1 T2 : lmodType R) (h : {linear T1 -> T2}) (v : tpower m T1),
    map_tpower h (f T1 v) = f T2 (map_tpower h v).

Lemma map_tpower_linear m (T1 T2 : lmodType R) (f : {linear T1 -> T2}) :
  linear (map_tpower (m:=m) f).
Proof. move=> x y z /=; apply/ffunP => vi; by rewrite !ffunE !linearE. Qed.
Canonical map_tpower_lin m T1 T2 f := Linear (@map_tpower_linear m T1 T2 f).

Lemma map_tpower_comp m (T1 T2 T3 : lmodType R) (f : T2 -> T3) (g : T1 -> T2) :
  map_tpower (m:=m) (f \o g) =1 map_tpower f \o map_tpower g.
Proof. by move=> v; apply/ffunP => vi; rewrite !ffunE. Qed.

Lemma map_tpcastE T m n (H : n = n) v :
  map_tpower (m:=m) (tpcast (T:=T) H) v = v.
Proof. by apply/ffunP => w /=; rewrite !ffunE tpcastE. Qed.

Lemma tpcastN m n (H : n = m) : naturality (mor_tpcast H).
Proof. move=> T1 T2 h v; apply/ffunP => vi; by rewrite !ffunE. Qed.

Definition eq_mor m n (f1 f2 : mor m n) := forall T : lmodType R, f1 T =1 f2 T.
Notation "f1 =e f2" := (eq_mor f1 f2).

Definition tsmor_fun m n (M : tmatrix n m) : morfun m n :=
  fun T v =>
    [ffun vi : n.-tuple I => \sum_(vj : m.-tuple I) (M vi vj : R) *: v vj].

Lemma tsmor_is_linear m n M T : linear (@tsmor_fun m n M T).
Proof.
move=> /= x y z; apply/ffunP => /= vi; rewrite !ffunE.
rewrite scaler_sumr -big_split; apply eq_bigr => /= vj _.
by rewrite !ffunE scalerDr !scalerA mulrC.
Qed.

Canonical tsmorfun m n (M : tmatrix n m) : mor m n :=
  fun T => Linear (@tsmor_is_linear m n M T).
Definition tsmor m n (M : tmatrix n m) : mor m n :=
  locked (tsmorfun M).

Lemma tsmorE m n (M : tmatrix n m) T v vi :
  tsmor M T v vi = \sum_(vj : m.-tuple I) (M vi vj : R) *: v vj.
Proof. by rewrite /tsmor -lock !ffunE. Qed.

Definition tpbasis m (vi : m.-tuple I) : tpower m R^o :=
  [ffun vj => (vi == vj)%:R].

Lemma tpbasisC m (vi vj : m.-tuple I) : tpbasis vi vj = tpbasis vj vi.
Proof. by rewrite !ffunE eq_sym. Qed.

Lemma sum_tpbasisK n (T : lmodType R) (vi : n.-tuple I) (F : tpower n T) :
  (\sum_vj (tpbasis vi vj *: F vj) = F vi).
Proof.
rewrite (bigD1 vi) //= !ffunE eqxx big1 ?(addr0,scale1r) //.
move=> vk; rewrite !ffunE eq_sym => /negbTE ->; by rewrite scale0r.
Qed.

Lemma sum_muleqr (A : finType) (S : comRingType) (F : A -> S) (v : A) :
  \sum_a F a * (v == a)%:R = F v.
Proof.
rewrite (bigD1 v) //= big1 ?(addr0,eqxx,mulr1) // => a av.
by rewrite eq_sym (negbTE av) mulr0.
Qed.

Lemma sum_tpbasisKo n (vi : n.-tuple I) (F : tpower n R) :
  (\sum_vj (F vj *: tpbasis vi vj) = F vi).
Proof. under eq_bigr do rewrite !ffunE. by rewrite sum_muleqr. Qed.

Lemma decompose_tpower m (T : lmodType R) (v : tpower m T) :
  v = (\sum_i map_tpower ( *:%R^~ (v i)) (tpbasis i)).
Proof.
apply/ffunP => vi; rewrite sum_ffunE -[LHS]sum_tpbasisK /=.
by apply eq_bigr => vj _; rewrite [RHS]ffunE tpbasisC.
Qed.

Definition morts m n (f : mor m n) : tmatrix n m :=
  [ffun vi => [ffun vj => f _ (tpbasis vj) vi]].

Lemma morts_eq m n (f g : mor m n) : f =e g -> morts f = morts g.
Proof. by move=> fg; apply/ffunP=>vi; apply/ffunP=>vj; rewrite !ffunE fg. Qed.

Lemma tsmorK m n : cancel (@tsmor m n) (@morts m n).
Proof.
move=> M; apply/ffunP => vi; apply/ffunP=> vj.
by rewrite !ffunE !tsmorE sum_tpbasisKo.
Qed.

Lemma mortsK n m (f : mor m n) :
  naturality f -> tsmor (morts f) =e f.
Proof.
move=> Hf T v.
rewrite [in RHS](decompose_tpower v) linear_sum.
apply/ffunP => /= vi; rewrite tsmorE !ffunE sum_ffunE /=.
apply eq_bigr => /= vj _; rewrite !ffunE.
set h : R^o -> T := *:%R^~ _.
have hlin : linear h by move=> x y z; rewrite /h scalerDl !scalerA.
by rewrite -(Hf _ _ (Linear hlin) (tpbasis vj)) ffunE.
Qed.

Lemma naturalityP m n (f : mor m n) :
  naturality f <-> exists M, f =e tsmor M.
Proof.
split => [Hf | [M] HM].
- by exists (morts f) => v T; rewrite mortsK.
- move=> T1 T2 h /= v; apply/ffunP => /= vi.
  rewrite !HM !(tsmorE,ffunE) linear_sum; apply eq_bigr => vj _.
  by rewrite linearZ_LR !ffunE.
Qed.

Lemma tsmorN m n M : naturality (@tsmor m n M).
Proof. by apply/naturalityP; esplit. Qed.

Definition ket_bra m (ket : tpower m R^o) (bra : tpower m R^o) : tsquare m :=
  [ffun vi => ket vi *: bra].

Definition mults m n p (M1 : tmatrix n m) (M2 : tmatrix m p) : tmatrix n p :=
  [ffun vi => [ffun vj => \sum_vk M1 vi vk * M2 vk vj]].

Notation "M1 '*t' M2" := (mults M1 M2).

Lemma multsA m n p q (M1 : tmatrix n m) (M2 : tmatrix m p) (M3 : tmatrix p q) :
  (M1 *t M2) *t M3 = M1 *t (M2 *t M3).
Proof.
apply/ffunP => vi; apply/ffunP => vj; rewrite !ffunE.
under eq_bigr do rewrite !ffunE big_distrl /=.
rewrite exchange_big /=; apply eq_bigr => vk _.
by rewrite !ffunE big_distrr /=; apply eq_bigr => vl _; rewrite mulrA.
Qed.

Definition idts m : tsquare m := [ffun vi => tpbasis vi].
Definition idmor n : endo n := fun T => GRing.idfun_linear _.

Lemma idmorE n : idmor n =e tsmor (idts n).
Proof.
move=> T v; apply/ffunP => vi.
rewrite /idmor tsmorE.
under eq_bigr do rewrite ffunE; by rewrite sum_tpbasisK.
Qed.

Lemma idmorN n : naturality (idmor n).
Proof. done. Qed.

Lemma mor_tpcastE n (H : n = n) : mor_tpcast H =e idmor n.
Proof. by move=> T v; rewrite /mor_tpcast /= tpcastE. Qed.

Definition transpose_tsquare m (M : tsquare m) : tsquare m :=
  [ffun vi => [ffun vj => M vj vi]].
Lemma transpose_tsquare_involutive m : involutive (@transpose_tsquare m).
Proof. move=> M. apply/ffunP => vi. apply/ffunP => vj. by rewrite !ffunE. Qed.

(* Tensor product of tsquare matrices *)
Section tensor_tsquare.
Variables m n : nat.

Definition tensor_tsquare (M1 : tsquare m) (M2 : tsquare n) : tsquare (m + n) :=
  [ffun vi => [ffun vj =>
     M1 (extract (lens_left m n) vi) (extract (lens_left m n) vj) *
     M2 (extract (lens_right m n) vi) (extract (lens_right m n) vj)]].

Lemma tensor_linearl (M2 : tsquare n) : linear (tensor_tsquare ^~ M2).
Proof.
move=> x M M'. apply/ffunP => vi. apply/ffunP => vj.
by rewrite !ffunE /= mulrDl scalerA.
Qed.

Lemma tensor_linearr (M1 : tsquare m) : linear (tensor_tsquare M1).
Proof.
move=> x M M'. apply/ffunP => vi. apply/ffunP => vj.
by rewrite !ffunE /= mulrDr !scalerA (mulrC x) -scalerA.
Qed.
End tensor_tsquare.

Section curry.
Variables (T : lmodType R) (n m : nat) (l : lens n m).

Definition curry (st : tpower n T) : tpower m (tpower (n-m) T) :=
  [ffun v : m.-tuple I =>
   [ffun w : (n-m).-tuple I => st (merge_indices l v w)]].

Definition uncurry (st : tpower m (tpower (n-m) T)) : tpower n T :=
  [ffun v : n.-tuple I => st (extract l v) (extract (lothers l) v)].

Lemma uncurryK : cancel uncurry curry.
Proof.
move=> v; apply/ffunP => v1; apply/ffunP => v2.
by rewrite !ffunE extract_merge extract_lothers_merge.
Qed.

Lemma curryK : cancel curry uncurry.
Proof. move=> v; apply/ffunP => w; by rewrite !ffunE merge_indices_extract. Qed.

Lemma curry_is_linear : linear curry.
Proof. move=>x y z; apply/ffunP=>vi; apply/ffunP =>vj; by rewrite !ffunE. Qed.
Canonical curry_lin := Linear curry_is_linear.

Lemma uncurry_is_linear : linear uncurry.
Proof. move => x y z; apply/ffunP=> vi; by rewrite !ffunE. Qed.
Canonical uncurry_lin := Linear uncurry_is_linear.

(* Special cases of curry/uncurry *)
Definition curry0 (v : T) : tpower 0 T := [ffun _ => v].
Definition curryn0 : tpower n T -> tpower n (tpower 0 T) :=
  map_tpower curry0.
Definition uncurry0 (v : tpower 0 T) : T := v [tuple].

Lemma curryn0E :
  curryn0 = fun v => [ffun vi => [ffun _ => v vi]].
Proof. reflexivity. Qed.

Lemma curry0_is_linear : linear curry0.
Proof. move=> x y z. apply/ffunP => vi. by rewrite !ffunE. Qed.
Lemma curryn0_is_linear : linear curryn0.
Proof. move=> x y z. apply/ffunP=> vi. apply/ffunP=> vj. by rewrite !ffunE. Qed.
Lemma uncurry0_is_linear : linear uncurry0.
Proof. move=> x y z. by rewrite /uncurry0 !ffunE. Qed.
Canonical curry0_lin := Linear curry0_is_linear.
Canonical curryn0_lin := Linear curryn0_is_linear.
Canonical uncurry0_lin := Linear uncurry0_is_linear.
End curry.

Section inner_prod_coprod.
Variable n : nat.
Let cast_uncurry T := map_tpower (m:=n) (tpcast (T:=T) (esym (addKn n n))).
Definition M_inner_prod (M : tsquare n) :=
  tsmor (curry0 (uncurry (lens_left n n) (cast_uncurry M))).
Definition M_inner_coprod (M : tsquare n) :=
  tsmor (curryn0 (uncurry (lens_left n n) (cast_uncurry M))).
Definition inner_prod : mor (n+n) 0 := M_inner_prod (idts _).
Definition inner_coprod : mor 0 (n+n) := M_inner_coprod (idts _).
End inner_prod_coprod.

Section tpaux.
Variables (k : nat) (T : lmodType R).

Definition tpall (v : T) : tpower k T := [ffun => v].
Lemma tpall_linear : linear tpall.
Proof. move=> a x y; apply/ffunP => i; by rewrite !ffunE. Qed.
Canonical tpall_lin := Linear tpall_linear.

Definition tpsum (v : tpower k T) : T := \sum_i v i.
Lemma tpsum_linear : linear tpsum.
Proof.
rewrite/tpsum => a x y /=. rewrite scaler_sumr -big_split /=.
apply eq_bigr=> i _; by rewrite !ffunE.
Qed.
Canonical tpsum_lin := Linear tpsum_linear.

Variable vi : k.-tuple I.
Definition tpsingle (v : T) : tpower k T :=
  [ffun vj => (vi == vj)%:R *: v].
Lemma tpsingle_linear : linear tpsingle.
Proof. move=> a x y; apply/ffunP => i; by rewrite !ffunE /= linearP. Qed.
Canonical tpsingle_lin := Linear tpsingle_linear.

Definition tpsel (v : tpower k T) := v vi.
Lemma tpsel_is_linear : linear tpsel.
Proof. by move=> x y z; rewrite /tpsel !ffunE. Qed.
Canonical tpsel_lin := Linear tpsel_is_linear.
End tpaux.

Section partial_trace.
Variables (n m : nat) (l : lens n m) (f : endo n).

Definition ptracefun (T : lmodType R) (v : tpower m T) : tpower m T :=
  \sum_(vi : (n-m).-tuple I)
    (map_tpower (tpsel vi) \o curry l \o f T
       \o uncurry l \o map_tpower (tpsingle vi)) v.

Lemma ptrace_is_linear T : linear (@ptracefun T).
Proof.
move=> a x y; rewrite /ptracefun !linear_sum -big_split /=.
apply eq_bigr => vi _; by rewrite !linearP.
Qed.

Definition ptrace : endo m := fun T => Linear (@ptrace_is_linear T).

Definition antifocus T := [linear of curry l \o f T \o uncurry l].
(* Lemma antifocus_naturality :
  naturality f -> naturality antifocus. *)

Lemma uncurry_tpsingle_naturality vi :
 naturality (fun T => [linear of uncurry (T:=T) l \o map_tpower (tpsingle vi)]).
Proof. by move=> T1 T2 h v; apply/ffunP => i; rewrite !ffunE linearE. Qed.

Lemma tpsel_curry_naturality vi :
  naturality (fun T => [linear of map_tpower (tpsel vi) \o curry (T:=T) l]).
Proof. by move=> T1 T2 h v; apply/ffunP => i; rewrite /tpsel !ffunE. Qed.

Lemma ptrace_naturality : naturality f -> naturality ptrace.
Proof.
move=> Nf T1 T2 h v.
rewrite /= /ptracefun linear_sum.
apply eq_bigr => vi _ /=.
move: (tpsel_curry_naturality vi h) => /= ->.
rewrite Nf.
by move: (uncurry_tpsingle_naturality vi h v) => ->.
Qed.
End partial_trace.

Section focus.
Variables (n m : nat) (l : lens n m) (tr : endo m).
Definition focus_fun : endofun n :=
  fun T (v : tpower n T) => uncurry l (tr _ (curry l v)).

Lemma focus_is_linear T : linear (@focus_fun T).
Proof.
move=> x y z.
apply/ffunP => vi; rewrite !ffunE.
have -> : curry l (T := T) = Linear (curry_is_linear l (T:=T)) by [].
by rewrite !linearP !ffunE.
Qed.

Definition focus : endo n := locked (fun T => Linear (@focus_is_linear T)).

Lemma focusE : focus = fun T => Linear (@focus_is_linear T).
Proof. by rewrite /focus; unlock. Qed.
End focus.

Lemma focus_naturality n m l tr : naturality tr -> naturality (@focus n m l tr).
Proof.
case/naturalityP => M /= NM; apply/naturalityP.
exists (morts (focus l (tsmor M))).
move=> T /= v; apply/ffunP => /= vi.
rewrite tsmorE !focusE !ffunE NM tsmorE sum_ffunE.
under [RHS]eq_bigr do rewrite !ffunE tsmorE sum_ffunE scaler_suml.
rewrite exchange_big /=; apply eq_bigr => vj _.
rewrite [in LHS](decompose_tpower v) !ffunE sum_ffunE scaler_sumr.
by apply eq_bigr => i _; rewrite !ffunE !scalerA.
Qed.

Section asym_focus.
Variables (n m p : nat) (l : lens (m+n) m) (l' : lens (p+n) p) (tr : mor m p).

Lemma addKn_any : (m + n - m = p + n - p)%N.
Proof. by rewrite !addKn. Qed.

Definition asym_focus_fun : morfun (m + n) (p + n) :=
  fun T (v : tpower (m + n) T) =>
    uncurry l' (map_tpower (tpcast addKn_any) (tr _ (curry l v))).

Lemma asym_focus_is_linear T : linear (@asym_focus_fun T).
Proof.
move=> x y z.
apply/ffunP => vi. rewrite !ffunE.
have -> : curry l (T := T) = Linear (curry_is_linear l (T:=T)) by [].
by rewrite !linearP !ffunE.
Qed.

Definition asym_focus : mor (m + n) (p + n) :=
  fun T => Linear (@asym_focus_is_linear T).
End asym_focus.

Lemma asym_focus_naturality n m p l l' tr :
  naturality tr -> naturality (@asym_focus n m p l l' tr).
Proof.
case/naturalityP => M /= NM; apply/naturalityP.
exists (morts (asym_focus l l' (tsmor M))).
move=> T /= v; apply/ffunP => /= vi; rewrite tsmorE !ffunE NM tsmorE sum_ffunE.
under [RHS]eq_bigr do rewrite !ffunE tsmorE sum_ffunE scaler_suml.
rewrite exchange_big /=; apply eq_bigr => vj _.
rewrite [in LHS](decompose_tpower v) !ffunE sum_ffunE scaler_sumr.
by apply eq_bigr => i _; rewrite !ffunE !scalerA.
Qed.

Lemma asym_focus_sym (m n : nat) (l : lens (m+n) m) (f : mor m m) :
  asym_focus l l f =e focus l f.
Proof.
move=> T v /=; rewrite focusE /= /asym_focus_fun /focus_fun /=.
by rewrite map_tpcastE.
Qed.

Section focus_props.
Variables (n m p : nat) (l : lens n m).

(* Identity *)
Lemma focusI tr : naturality tr -> focus (lens_id n) tr =e tr.
Proof.
rewrite focusE => /naturalityP [f Hf] /= T v.
apply/ffunP => /= vi.
rewrite /focus_fun !{}Hf {tr} !ffunE !tsmorE sum_ffunE.
apply eq_bigr => vj _; rewrite !ffunE extract_lens_id.
congr (_ *: v _).
apply eq_from_tnth => i; by rewrite tnth_mktuple index_lens_id -tnth_nth.
Qed.

(* Equality *)
Lemma focus_eq (f1 f2 : endo m) : f1 =e f2 -> focus l f1 =e focus l f2.
Proof. move=> Heq T v /=; by rewrite 2!focusE /= /focus_fun Heq. Qed.

(* Vertical composition of morphisms *)
Section comp_mor.
Variables (r q s : nat) (tr : mor q s) (tr' : mor r q).
Definition comp_mor : mor r s := fun A => GRing.comp_linear (tr A) (tr' A).

Lemma comp_naturality : naturality tr -> naturality tr' -> naturality comp_mor.
Proof. move=> N1 N2 T1 T2 f v; by rewrite N1 N2. Qed.
End comp_mor.
Notation "f \v g" := (comp_mor f g).

Definition cast_mor m2 n2 m1 n1 (f : mor m1 n1)
           (Hm : m1 = m2) (Hn : n1 = n2) : mor m2 n2 :=
  mor_tpcast Hn \v f \v mor_tpcast (esym Hm).

Section comp_mor_facts.
Variables (q : nat) (f : mor m n) (g : mor n p) (h : mor p q).

Lemma comp_morA : h \v (g \v f) =e (h \v g) \v f.
Proof. by []. Qed.

Lemma comp_morE T v : (g \v f) T v = g T (f T v).
Proof. by []. Qed.
End comp_mor_facts.

Lemma focus_comp r q (tr tr' : endo q) (lq : lens r q) :
  focus lq (tr \v tr') =e focus lq tr \v focus lq tr'.
Proof.
move=> T v; apply/ffunP => /= vi.
by rewrite !focusE /focus_fun /= uncurryK.
Qed.

Lemma tsmor_comp (M : tmatrix n m)  (N : tmatrix m p) :
  tsmor (mults M N) =e tsmor M \v tsmor N.
Proof.
move=> T v; apply/ffunP => vi; rewrite !tsmorE.
under eq_bigr do rewrite !ffunE !scaler_suml.
rewrite exchange_big /=.
apply eq_bigr => vk _; rewrite tsmorE !(scaler_suml,scaler_sumr).
by apply eq_bigr => vj _; rewrite scalerA.
Qed.

(* Horizontal composition of endomorphisms *)
Lemma focusC (l' : lens n p) tr tr' :
  [disjoint l & l'] -> naturality tr -> naturality tr' ->
  focus l tr \v focus l' tr' =e focus l' tr' \v focus l tr.
Proof.
rewrite !focusE => Hdisj /naturalityP [f Hf] /naturalityP [f' Hf'] T v /=.
apply/ffunP => /= vi.
rewrite /focus_fun !{}Hf !{}Hf' {tr tr'} !ffunE !tsmorE !sum_ffunE.
under eq_bigr do rewrite !ffunE tsmorE !sum_ffunE scaler_sumr.
rewrite exchange_big; apply eq_bigr => /= vj _.
rewrite !ffunE tsmorE !sum_ffunE scaler_sumr; apply eq_bigr => /= vk _.
rewrite !ffunE !scalerA [in RHS]mulrC.
congr (f _ vk * f' _ vj *: v _).
- by rewrite extract_merge_disjoint // disjoint_sym.
- by rewrite extract_merge_disjoint.
- by rewrite !merge_indices_extract_others inject_disjointC.
Qed.

Lemma focus_tensor (M : tsquare m) (M' : tsquare n) :
  focus (lens_left m n) (tsmor M) \v focus (lens_right m n) (tsmor M') =e
  tsmor (tensor_tsquare M M').
Proof.
move=> T v; apply/ffunP => /= vi.
rewrite focusE !(ffunE,tsmorE) !sum_ffunE.
under eq_bigr do rewrite !focusE !(ffunE,tsmorE) !sum_ffunE scaler_sumr.
rewrite reindex_left_right.
apply eq_bigr => /= vj _; rewrite !ffunE !merge_indices_extract_others.
rewrite extract_inject; last by rewrite disjoint_sym lens_left_right_disjoint.
by rewrite scalerA inject_all // lens_left_right_disjoint.
Qed.

(* Associativity of actions of lenses *)
Lemma focusM (l' : lens m p) tr : naturality tr ->
  focus (lens_comp l l') tr =e focus l (focus l' tr).
Proof.
case/naturalityP => f Hf T v.
rewrite !focusE /focus_fun /= !{}Hf {tr}.
apply/ffunP => /= vi.
rewrite !ffunE !tsmorE (extract_lothers_comp dI) -!extract_comp.
rewrite -[in RHS]lothers_in_l_comp -(lothers_notin_l_comp l l') !sum_ffunE.
apply eq_bigr => /= vj _; rewrite !ffunE.
congr (_ *: v _).
exact: merge_indices_comp.
Qed.

(* Variant for disjoint lenses, used in unitary.v *)
Variable T : lmodType R.
Lemma focus_others (l' : lens (n-m) p) (f : endo p) (t : tpower n T) :
  naturality f ->
  focus (lens_comp (lothers l) l') f T t =
  uncurry l (map_tpower (m:=m) (focus l' f T) (curry l t)).
  (* parametricity prevents writing it this way:
     focus l (fun _ => Linear (map_tpower_linear (focus l' f T))) T t. *)
Proof.
case/naturalityP => M Hf; apply/ffunP => vi.
rewrite !focusE !ffunE /= -!extract_comp !Hf !tsmorE /= !sum_ffunE.
apply eq_bigr => vj _; by rewrite !ffunE merge_indices_comp_others.
Qed.
End focus_props.
Notation "f \v g" := (comp_mor f g).
Notation tsapp l M := (focus l (tsmor M)).

(* Too complicated
Lemma asym_focusC n m p n' m' p' (l1 : lens (m+n) m) (l2 : lens (p+n) p)
      (l3 : lens (m'+n') m') (l4 : lens (p'+n') p') (tr : mor m p)
      (tr' : mor m' p') :
  [disjoint map val l2 & map val l3] -> p + n = m' + n' ->
  asym_focus l3 l4 tr' \v asym_focus l1 l2 tr \v =e ???
*)

Lemma asym_focusC n m p (l1 : lens (m+n) m) (l2 : lens (p+n) p)
      (g : mor m p) (f : endo n) :
  naturality f -> naturality g ->
  focus (cast_lens (lothers l2) (addKn _ _)) f \v asym_focus l1 l2 g =e
  asym_focus l1 l2 g \v focus (cast_lens (lothers l1) (addKn _ _)) f.
Proof.
rewrite !focusE => /naturalityP [Mf Hf] /naturalityP [Mg Hg] T v /=.
apply/ffunP => /= vi.
rewrite /focus_fun /asym_focus_fun !{}Hf {f} !{}Hg {g}.
rewrite !ffunE !tsmorE !sum_ffunE.
under eq_bigr do rewrite !ffunE tsmorE !sum_ffunE scaler_sumr.
rewrite exchange_big; apply eq_bigr => /= vj _.
rewrite !ffunE tsmorE !sum_ffunE scaler_sumr; apply eq_bigr => /= vk _.
rewrite !ffunE !scalerA [in RHS]mulrC.
congr (Mf _ vk * Mg _ vj *: v _).
- apply val_inj => /=.
  set w := cast_tuple _ _.
  by move/(f_equal val): (extract_lothers_merge dI l1 vj w) => /= ->.
- rewrite extract_merge_disjoint //.
  apply/pred0P => /= i.
  rewrite simpl_predE /= andbC /=.
  case Hi: (i \in l2) => //=.
  by rewrite mem_lensE /= mem_others Hi.
- rewrite !merge_indices_extract_others.
  apply eq_from_tnth => i /=.
  rewrite !tnth_mktuple /=.
  case/boolP: (i \in l1) => Hil1.
    rewrite !nth_lens_index.
    rewrite nth_default // memNindex ?(mem_others,Hil1) //.
    by rewrite (eqP (size_others l1)) size_tuple addKn.
  rewrite !nth_lens_out //.
  move/negbF: Hil1.
  rewrite -mem_others => /negbFE.
  have -> : others l1 = cast_lens (lothers l1) (addKn _ _) by [].
  move=> Hil1.
  rewrite !nth_lens_index (tnth_nth dI).
  congr nth.
  have -> : others l2 = cast_lens (lothers l2) (addKn _ _) by [].
  transitivity (map_tuple (tnth (inject (cast_lens (lothers l2) (addKn p n))
                           vi vk)) (cast_lens (lothers l2) (addKn p n))) => //.
  apply f_equal, eq_from_tnth => j.
  rewrite tnth_map /= tnth_mktuple nth_lens_index ?mem_tnth // => H.
  rewrite -(nth_lens_index H dI) nthK /=.
  + by rewrite -?tnth_nth.
  + by rewrite uniq_others.
  + by rewrite (eqP (size_others _)) addKn inE.
Qed.

Lemma focus_tensor' n m p (l : lens n m) (l' : lens n p) (H : [disjoint l & l'])
      (M : tsquare m) (M' : tsquare p) :
  tsapp l M \v tsapp l' M' =e tsapp (lens_cat H) (tensor_tsquare M M').
Proof.
rewrite {1}(lens_comp_right H) {1}(lens_comp_left H) => T v /=.
rewrite focusM; last by apply/tsmorN.
rewrite -> focusM; last by apply/tsmorN.
have /= <- := focus_comp _ _ _ v.
move: T v; exact/focus_eq/focus_tensor.
Qed.

(*
Section narrow.
Variables (n m : nat) (l : lens n m).
Definition narrow_in (T : lmodType R) (st : tpower m T) : tpower n T :=
  [ffun v : n.-tuple I => st (extract l v)].
Definition narrow_out (T : lmodType R) (st : tpower n T) : tpower m T :=
  [ffun v : m.-tuple I => st (inject l [tuple of nseq n dI] v)].
Definition narrow (f : endofun n) : endofun m :=
  fun T (v : tpower m T) => narrow_out (f T (narrow_in v)).
End narrow.
Lemma narrow_focus n m p (l : lens n m) (l' : lens n p)
      (H : [disjoint l & l']) f (T : lmodType R) v :
  narrow (lothers l) (focus l' f) v =
  focus (lmake_comp H) f T v.
Proof.
apply/ffunP => vi.
rewrite /narrow !focusE /= !ffunE.
Abort.
*)
End tensor_space.

Notation "f1 =e f2" := (eq_mor f1 f2).
Notation "f \v g" := (comp_mor f g).
Notation "M1 '*t' M2" := (mults M1 M2).
Notation tsapp l M := (focus l (tsmor M)).

(* Conversion between tpower and vector space *)

Section index_of_vec_bij.
Local Open Scope nat_scope.
Variable I : finType.
Variable dI : I.
Let vsz m := #|I| ^ m.

Fixpoint index_of_vec_rec (v : seq I) : nat :=
  match v with
  | nil => 0
  | i :: v' => enum_rank i + #|I| * index_of_vec_rec v'
  end.

Lemma index_of_vec_ltn m (v : seq I) :
  size v = m -> index_of_vec_rec v < vsz m.
Proof.
rewrite /vsz.
elim: v m => [|i v IH []] //= m.
  move <-; by rewrite expn0.
case=> Hm; rewrite expnS.
case: enum_rank => j /= Hj.
have : #|I| ^ m > 0.
  rewrite -(expn0 #|I|) leq_pexp2l //.
  by case: #|I| Hj.
move CI: (#|I| ^ m) => [|sz] // _.
by rewrite mulnS -addSn leq_add // leq_mul // -ltnS -CI IH.
Qed.

Definition index_of_vec m (v : m.-tuple I) : 'I_(vsz m).
exists (index_of_vec_rec (rev v)).
abstract (by rewrite index_of_vec_ltn // size_rev size_tuple).
Defined.

Lemma card_inh : #|I| > 0.
Proof. by rewrite -cardsT card_gt0; apply/set0Pn; exists dI; rewrite inE. Qed.

Fixpoint vec_of_index_rec (m i : nat) : seq I :=
  match m with
  | 0 => nil
  | m.+1 =>
    enum_val (Ordinal (ltn_pmod i card_inh)) :: vec_of_index_rec m (i %/ #|I|)
  end.

Lemma vec_of_index_size m i : size (vec_of_index_rec m i) = m.
Proof. by elim: m i => // m IH [|i] /=; rewrite IH. Qed.

Definition vec_of_index m (i : 'I_(vsz m)) : m.-tuple I.
exists (rev (vec_of_index_rec m i)).
abstract (by case: i => i /= _; rewrite size_rev vec_of_index_size).
Defined.

Lemma vec_of_index_recK m i :
  i < vsz m -> index_of_vec_rec (vec_of_index_rec m i) = i.
Proof.
rewrite /vsz.
elim: m i => [|m IH /= i Hi]; first by case; rewrite expn0 // ltnS.
rewrite enum_valK IH /=; first by rewrite addnC mulnC -divn_eq.
rewrite -(ltn_pmul2r card_inh) (leq_ltn_trans (leq_trunc_div _ _)) //.
by rewrite mulnC -expnS.
Qed.

Lemma vec_of_indexK m : cancel (@vec_of_index m) (@index_of_vec m).
Proof.
rewrite /index_of_vec /vec_of_index /= => -[i] Hi.
apply val_inj; by rewrite /= revK vec_of_index_recK.
Qed.

Lemma index_of_vecK m : cancel (@index_of_vec m) (@vec_of_index m).
Proof.
rewrite /index_of_vec /vec_of_index => -[t Ht].
apply/val_inj => /=.
rewrite -[RHS]revK; congr rev.
move/eqP: Ht; rewrite -size_rev.
elim: (rev t) m => {t} [|i t IH] m <- //=.
congr (_ :: _).
  rewrite (_ : Ordinal _ = enum_rank i) ?enum_rankK //.
  apply val_inj => /=.
  by rewrite addnC mulnC modnMDl modn_small.
rewrite divnDr.
  by rewrite divn_small // add0n mulKn ?card_inh // IH.
exact/dvdn_mulr/dvdnn.
Qed.

Lemma index_of_vec_bij m : bijective (@index_of_vec m).
Proof.
exists (@vec_of_index m); [exact: index_of_vecK | exact: vec_of_indexK].
Qed.

Lemma vec_of_index_bij m : bijective (@vec_of_index m).
Proof.
exists (@index_of_vec m); [exact: vec_of_indexK | exact: index_of_vecK].
Qed.
End index_of_vec_bij.

(* tpower n R^o forms a vector space of size #|I|^m *)
Section vector.
Variable (I : finType) (R : comRingType) (dI : I).
Let vsz m := (#|I| ^ m)%N.
Let tmatrix := tmatrix I R.

Section mxtmatrix.
Variables m n : nat.
Definition mxtmatrix (M : 'M[R]_(vsz m,vsz n)) : tmatrix m n :=
  [ffun vi => [ffun vj => M (index_of_vec vi) (index_of_vec vj)]].

Definition tmatrixmx (M : tmatrix m n) : 'M[R]_(vsz m,vsz n) :=
  \matrix_(i,j) M (vec_of_index dI i) (vec_of_index dI j).

Lemma tmatrixmxK : cancel tmatrixmx mxtmatrix.
Proof.
move=> v; apply/ffunP => vi; apply/ffunP => vj.
by rewrite !ffunE mxE !index_of_vecK.
Qed.

Lemma mxtmatrixK : cancel mxtmatrix tmatrixmx.
Proof.
move=> v; apply/matrixP => i j; by rewrite mxE !ffunE !vec_of_indexK.
Qed.
End mxtmatrix.

Lemma tmatrixmx_mul m n p (M1 : tmatrix m n) (M2 : tmatrix n p) :
  tmatrixmx (M1 *t M2) = tmatrixmx M1 *m tmatrixmx M2.
Proof.
apply/matrixP => i j; rewrite !mxE !ffunE.
rewrite (reindex (@index_of_vec I n)) /=.
  apply eq_bigr => vi _; by rewrite !mxE index_of_vecK.
exists (@vec_of_index _ dI n) => x y; by rewrite (vec_of_indexK,index_of_vecK).
Qed.

Lemma mxtmatrix_mul m n p (M1 : 'M_(vsz m,vsz n)) (M2 : 'M_(vsz n,vsz p)) :
  mxtmatrix (M1 *m M2) = mxtmatrix M1 *t mxtmatrix M2.
Proof.
apply/ffunP => vi; apply/ffunP => vj; rewrite !ffunE !mxE.
rewrite (reindex (@index_of_vec I n)) /=.
  apply eq_bigr => vk _; by rewrite !ffunE.
exists (@vec_of_index _ dI n) => x y; by rewrite (vec_of_indexK,index_of_vecK).
Qed.

Lemma tmatrixmx_id m : tmatrixmx (idts I R m) = (1%:M).
Proof.
apply/matrixP => i j; rewrite !mxE !ffunE.
by rewrite (inj_eq (bij_inj (vec_of_index_bij dI m))).
Qed.

Lemma mxtmatrix_id m : mxtmatrix (1%:M) = idts I R m.
Proof.
apply/ffunP => vi; apply/ffunP => vj; rewrite !ffunE mxE.
by rewrite (inj_eq (bij_inj (index_of_vec_bij dI m))).
Qed.

Lemma mul1ts m n (M : tmatrix  m n) : idts I R m *t M = M.
Proof.
by rewrite -[LHS]tmatrixmxK tmatrixmx_mul tmatrixmx_id mul1mx tmatrixmxK.
Qed.
Lemma mults1 m n (M : tmatrix  m n) : M *t idts I R n = M.
Proof.
by rewrite -[LHS]tmatrixmxK tmatrixmx_mul tmatrixmx_id mulmx1 tmatrixmxK.
Qed.

Definition vec_tpower m (X : 'rV[R]_(vsz m)) : tpower I m R^o :=
  [ffun vi => X ord0 (index_of_vec vi)].

Definition tpower_vec m (X : tpower I m R^o) : 'rV[R]_(vsz m) :=
  \row_i X (vec_of_index dI i).

Definition mxmor m n (M : 'M_(vsz m,vsz n)) := tsmor (mxtmatrix M).

Lemma tpower_vector m : Vector.axiom (vsz m) (tpower I m R^o).
Proof.
exists (@tpower_vec m).
- move=> x /= y z. apply/rowP => i. by rewrite !(ffunE,mxE).
- exists (@vec_tpower m).
  + move=> v. apply/ffunP => vi. by rewrite !(ffunE,mxE) index_of_vecK.
  + move=> X. apply/rowP => i. by rewrite !(ffunE,mxE) vec_of_indexK.
Qed.
End vector.

(* Helper lemmas for computation *)
Section enum_indices.
Variable I : finType.
Variable enumI : seq I.
Hypothesis uniq_enumI : uniq enumI.
Hypothesis mem_enumI : forall i, i \in enumI.

Fixpoint enum_indices m : seq (m.-tuple I) :=
  match m with
  | 0 => [:: [tuple of [::]]]
  | S m =>
    allpairs (fun x (t : m.-tuple _) => [tuple of x :: val t])
             enumI (enum_indices m)
  end.

Lemma mem_enum_indices m t : t \in enum_indices m.
Proof.
elim: m t => [|m IH] [[|i t] Hlen] //=.
apply/flatten_mapP.
exists i => //.
case/eqP: (Hlen) => /eqP Hlen'.
apply/mapP; exists (Tuple Hlen') => //; exact/val_inj.
Qed.

Lemma size_enum_indices m : size (enum_indices m) = (size enumI ^ m)%N.
Proof. elim: m => //= m IH; by rewrite size_allpairs IH expnS. Qed.

Lemma uniq_enum_indices m : uniq (enum_indices m).
Proof.
rewrite /is_true -(enum_uniq (tuple_finType m I)).
apply eq_uniq.
  rewrite -cardT card_tuple size_enum_indices; congr expn.
  move/card_uniqP: uniq_enumI => <-.
  apply eq_card => i; by rewrite mem_enumI.
move=> t. by rewrite mem_enum_indices mem_enum.
Qed.

Lemma forall_indicesP n (P : pred (n.-tuple I)) :
  reflect (forall v, P v) (all P (enum_indices n)).
Proof.
apply (iffP idP).
  move=> H vi.
  have : vi \in enum_indices _ by rewrite mem_enum_indices.
  by apply/allP: vi.
move=> H; by apply/allP.
Qed.

Lemma eq_from_indicesP n (T : eqType) (v w : tpower I n T) :
  reflect (v = w) (all (fun x => v x == w x) (enum_indices n)).
Proof.
apply/(iffP (forall_indicesP _)) => [H | -> //].
apply/ffunP => vi; exact/eqP.
Qed.

Lemma sum_enum_indices (L : zmodType) m (F : m.-tuple I -> L) :
  (\sum_vi F vi = foldr +%R 0 (map F (enum_indices m))).
Proof.
rewrite foldrE big_map [RHS]big_uniq ?uniq_enum_indices //=.
apply/esym/eq_bigl => vi. exact/mem_enum_indices.
Qed.
End enum_indices.

Section enum2.
Let I := [finType of 'I_2].

Definition enum2 : seq I := [:: 0%:O; 1%:O].
Lemma uniq_enum2 : uniq enum2. Proof. by []. Qed.
Lemma mem_enum2 i : i \in enum2.
Proof. by rewrite !inE; case: i => -[|[]]. Qed.
End enum2.