Set Warnings "-notation-overridden,-parsing,-deprecated-hint-without-locality".
From Coq Require Import Strings.String.  (* for manual grading *)
From Coq Require Export Bool.Bool.
From Coq Require Export Arith.Arith.
From Coq Require Export Arith.EqNat.
From Coq Require Export Lia.
From Coq Require Export Lists.List.
Export ListNotations.
From Coq Require Export Permutation.

Print reflect.

Notation  "a >=? b" := (Nat.leb b a)
                          (at level 70) : nat_scope.
Notation  "a >? b"  := (Nat.ltb b a)
                         (at level 70) : nat_scope.

Lemma eqb_reflect : forall x y, reflect (x = y) (x =? y).
Proof.
  intros x y. apply iff_reflect. symmetry.
  apply Nat.eqb_eq.
Qed.
Lemma ltb_reflect : forall x y, reflect (x < y) (x <? y).
Proof.
  intros x y. apply iff_reflect. symmetry.
  apply Nat.ltb_lt.
Qed.
Lemma leb_reflect : forall x y, reflect (x <= y) (x <=? y).
Proof.
  intros x y. apply iff_reflect. symmetry.
  apply Nat.leb_le.
Qed.

Example reflect_example1: forall a,
    (if a <? 5 then a else 2) < 6.
Proof.
  intros a.
  (* The next two lines aren't strictly necessary, but they
     help make it clear what destruct does. *)
  assert (R: reflect (a < 5) (a <? 5)) by apply ltb_reflect.
  remember (a <? 5) as guard.
  destruct R as [H|H] eqn:HR.
  * (* ReflectT *) lia.
  * (* ReflectF *) lia.
Qed.

Hint Resolve ltb_reflect leb_reflect eqb_reflect : bdestruct.

Ltac bdestruct X :=
  let H := fresh in let e := fresh "e" in
   evar (e: Prop);
   assert (H: reflect e X); subst e;
    [eauto with bdestruct
    | destruct H as [H|H];
       [ | try first [apply not_lt in H | apply not_le in H]]].

Example reflect_example2: forall a,
    (if a <? 5 then a else 2) < 6.
Proof.
  intros.
  bdestruct (a <? 5); (* instead of: destruct (ltb_reflect a 5). *)
  lia.
Qed.

Definition maybe_swap (al: list nat) : list nat :=
  match al with
  | a :: b :: ar => if a >? b then b :: a :: ar else a :: b :: ar
  | _            => al
  end.

Theorem maybe_swap_idempotent: forall al,
    maybe_swap (maybe_swap al) = maybe_swap al.
Proof.
  intros [ | a [ | b al]]; simpl; try reflexivity.
  bdestruct (a >? b); simpl.
  - bdestruct (b >? a); simpl.
    + lia.
    + reflexivity.
  - bdestruct (a >? b); simpl.
    + lia.
    + reflexivity.
Qed.

Print Permutation.
Search Permutation.

Example butterfly: forall b u t e r f l y : nat,
  Permutation ([b;u;t;t;e;r]++[f;l;y]) ([f;l;u;t;t;e;r]++[b;y]).
Proof.
  intros.
  change [b;u;t;t;e;r] with ([b]++[u;t;t;e;r]).
  change [f;l;u;t;t;e;r] with ([f;l]++[u;t;t;e;r]).
  remember [u;t;t;e;r] as utter. clear Hequtter.
  change [f;l;y] with ([f;l]++[y]).
  remember [f;l] as fl. clear Heqfl.
  replace ((fl ++ utter) ++ [b;y]) with (fl ++ utter ++ [b;y])
    by apply app_assoc.
  apply perm_trans with (fl ++ [y] ++ ([b] ++ utter)).
  - replace (fl ++ [y] ++ [b] ++ utter) with ((fl ++ [y]) ++ [b] ++ utter).
    + apply Permutation_app_comm.
    + rewrite <- app_assoc. reflexivity.
  -
    apply Permutation_app_head.
    apply perm_trans with (utter ++ [y] ++ [b]).
    + replace ([y] ++ [b] ++ utter) with (([y] ++ [b]) ++ utter).
      * apply Permutation_app_comm.
      * rewrite app_assoc. reflexivity.
    + apply Permutation_app_head.
      apply perm_swap.
Qed.

Check perm_skip.
Check perm_trans.
Check Permutation_refl.
Check Permutation_app_comm.
Check app_assoc.
Check app_nil_r.
Check app_comm_cons.
Example permut_example: forall (a b: list nat),
  Permutation (5 :: 6 :: a ++ b) ((5 :: b) ++ (6 :: a ++ [])).
Proof.
  intros.
  rewrite app_nil_r.
  rewrite <- app_comm_cons.
  apply perm_skip.
  rewrite app_comm_cons.
  apply Permutation_app_comm.
Qed.

Check Permutation_cons_inv.
Check Permutation_length_1_inv.
Example not_a_permutation:
  ~ Permutation [1;1] [1;2].
Proof.
  unfold not. intros.
  apply Permutation_cons_inv in H.
  apply Permutation_length_1_inv in H.
  discriminate H.
Qed.

Theorem maybe_swap_perm: forall al,
  Permutation al (maybe_swap al).
Proof.
  unfold maybe_swap.
  destruct al as [ | a [ | b al]].
  - simpl. apply perm_nil.
  - apply Permutation_refl.
  - bdestruct (b <? a).
    + apply perm_swap.
    + apply Permutation_refl.
Qed.

Definition first_le_second (al: list nat) : Prop :=
  match al with
  | a :: b :: _ => a <= b
  | _           => True
  end.

Theorem maybe_swap_correct: forall al,
    Permutation al (maybe_swap al) /\ first_le_second (maybe_swap al).
Proof.
  intros. split.
  - apply maybe_swap_perm.
  - unfold maybe_swap.
    destruct al as [ | a [ | b al]]; simpl; auto.
    bdestruct (a >? b); simpl; lia.
Qed.

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

Print Forall.

Theorem Forall_perm: forall {A} (f: A -> Prop) al bl,
  Permutation al bl ->
  Forall f al -> Forall f bl.
Proof.
  intros.
  induction H.
  - assumption.
  - inv H0.
    apply Forall_cons.
    + assumption.
    + apply (IHPermutation H4).
  - inv H0. inv H3.
    apply Forall_cons.
    assumption.
    apply Forall_cons.
    assumption.
    assumption.
  - apply IHPermutation2.
    apply IHPermutation1.
    assumption.
Qed.


(* The El Gamal encryption scheme and a proof that it is IND-CPA-secure. *)

Set Implicit Arguments.

Require Import FCF.
Require Import RndGrpElem.
Require Import Encryption_PK.
Require Import DiffieHellman.
Require Import OTP.

Local Open Scope group_scope.

Section ElGamal.

  Context`{FCG : FiniteCyclicGroup}.

  Hypothesis GroupElement_EqDec : EqDec GroupElement.

  Definition ElGamalKeygen :=
    m <-$ [0 .. order);
    ret (m, g^m).

  Definition ElGamalEncrypt(msg key : GroupElement) :=
    m <-$ [0 .. order);
    ret (g^m, key^m * msg).

  Variable A_State : Set.
  Hypothesis A_State_EqDec : EqDec A_State.

  Variable A1 : GroupElement -> Comp (GroupElement * GroupElement * A_State).
  Hypothesis wfA1 : forall x, well_formed_comp (A1 x).

  Variable A2 : (GroupElement * GroupElement * A_State) -> Comp bool.
  Hypothesis wfA2 : forall x, well_formed_comp (A2 x).

  (* Build an adversary from A1 and A2 that can win DDH *)
  Definition B(g_xyz : (GroupElement * GroupElement * GroupElement)) : Comp bool :=
    [gx, gy, gz] <-3 g_xyz;
    [p0, p1, s] <-$3 A1(gx);
    b <-$ {0,1};
    pb <- if b then p0 else p1;
      c <- (gy, gz * pb);
      b' <-$ (A2 (c, s));
      ret (eqb b b').

  Theorem ElGamal_IND_CPA0 :
    Pr[IND_CPA_G ElGamalKeygen ElGamalEncrypt A1 A2] ==
    Pr[(@DDH0 _ _ _ _ _ g order B)].

    unfold IND_CPA_G, DDH0, ElGamalKeygen, ElGamalEncrypt, B.

    inline_first.
    comp_skip.

    dist_at dist_inline rightc 1%nat.
    comp_swap rightc.
    comp_skip.

    destruct x0.
    destruct p.

    dist_at dist_inline rightc 1%nat.
    comp_swap rightc.
    comp_skip.

    comp_inline leftc.
    comp_skip.

    comp_inline rightc.
    comp_skip.
    rewrite groupExp_mult; intuition.

    comp_simp.
    intuition.
  Qed.

  Definition G1 :=
    gx <-$ RndG;
    gy <-$ RndG;
    [p0, p1, s] <-$3 (A1 gx);
    b <-$ {0, 1};
    gz' <-$ (
    pb <- if b then p0 else p1;
    gz <-$ RndG ; ret (gz * pb));
    b' <-$ (A2 (gy, gz', s));
    ret (eqb b b').

  Definition G2 :=
    gx <-$ RndG;
    gy <-$ RndG;
    [p0, p1, s] <-$3 (A1 gx);
    gz <-$ RndG ;
    b' <-$ (A2 (gy, gz, s));
    b <-$ {0, 1};
    ret (eqb b b').

  Theorem ElGamal_G1_DDH1 :
    Pr [ G1] == Pr [ (@DDH1 _ _ _ _ _ g order B) ].

    unfold G1, DDH1, B, RndGrpElem.

    inline_first.
    comp_skip.
    inline_first.
    comp_skip.

    comp_at comp_inline rightc 1%nat.
    comp_swap rightc.
    comp_skip.

    destruct x1.
    destruct p.

    dist_at dist_inline leftc 1%nat.
    dist_at dist_inline leftc 1%nat.
    comp_swap leftc.
    comp_skip.

    comp_inline rightc.
    comp_skip.

    inline_first.
    comp_skip.

    comp_simp.
    intuition.

  Qed.

  Theorem ElGamal_G1_G2 :
    Pr[G1] == Pr[G2].

    unfold G1, G2, B.

    (* we do this step in the program logic *)
    eapply comp_spec_eq_impl_eq.

    do 3 comp_skip.

    prog_at_r prog_swap 1%nat.
    prog_swap_r.
    comp_skip.

    eapply comp_spec_symm.
    eapply comp_spec_seq; eauto with inhabited.
    eapply group_OTP_r.
    intuition.
    subst.
    eapply comp_spec_consequence.
    eapply comp_spec_eq_refl.
    intuition.
  Qed.

  Theorem ElGamal_G2_OneHalf :
    Pr [G2] == 1 / 2.

    unfold G2.

    (* ignore the first 5 commands *)
    do 3 dist_irr_l.
    comp_simp.
    do 2 dist_irr_l.
    (* compute the probability *)
    dist_compute.
  Qed.

  Theorem ElGamal_IND_CPA_Advantage :
    (IND_CPA_Advantage ElGamalKeygen ElGamalEncrypt A1 A2) ==
    (@DDH_Advantage _ _ _ _ _ g order B).

    unfold IND_CPA_Advantage, DDH_Advantage.
    eapply ratDistance_eqRat_compat.
    eapply ElGamal_IND_CPA0.
    rewrite <- ElGamal_G1_DDH1.
    rewrite ElGamal_G1_G2.
    symmetry.
    eapply ElGamal_G2_OneHalf.
  Qed.

End ElGamal.


Require Export ConnectivityGraph.
Require Export Layouts.
Require Export MappingConstraints.
Require Import StandardGateSet.
Import StdList.

(* Alogrithm 3.1 *)

Fixpoint hierarchical_product_permuter_part1 pi permuter_g1g2 (deg_ham : nat) :=  (* line 2 *)
  match deg_ham with
  | 0 => pi
  | n =>
      let r = partial_permutation permuter_g1g2  (* line 1*)
          g' = route_v_to_k r permuter_g1g2  (* line 3,4 *)
          new_g = route_communicator_v_g1 g'  (* line 5,6 *)
      in
      hierarchical_product_permuter_part1 pi new_g (n-1)
  end.

Fixpoint hierarchical_product_permuter_part2 v1 pi permuter_g1g2 :=  (* line 7 *)
  match v1 with
  | [] => pi
  | (h :: t) =>
      let
        new_pi = route_v_Vi_to_pi h pi permuter_g1g2  (*line 8 *)
      in
      hierarchical_product_permuter_part2 t new_pi permuter_g1g2
  end,

Definition hierarchical_product_permuter pi permuter_g1g2 deg_ham := (* Algorithm 3.1 *)
let
  new_pi = hierarchical_product_permuter_part1 pi permuter_g1g2 deg_ham
  v1 = V(new_pi,1)
in
hierarchical_product_permuter_part2 v1 new_pi permuter_g1g2

(* Algorithm 3.2 *)

Fixpoint choose_distinct_sets pi sigma r rowi_V2 :=
match rowi_V2 with
| [] => sigma
| i :: tail => let v = (dom pi) - (dom sigma)
                   E = (head v, pi_v (v,1), current_v (v, i)) (* line 4 *)
                   U, V = list_n (num_of (V1 pi)) (* ni = |Vi| *) (* line 6 *)
                   G = (U, V, E) (* line 6 *)
                   V_match = v_set (find_minimum_weight G (add_u_v_e G r)) (* line 7~11 *)
                   new_sigma =  app sigma (v_list V_match) (* line 12 *)
               in
                   choose_distinct_sets pi new_sigma (r-1) t

(* Algorithm 3.3 *)

Fixpoint routing_tokens_to_destination pi :=
if (pi == id_dom pi)
then (*return seq of transpositions *)
else
  if exists_happy_swap
  then routing_tokens_to_destination (do_transposition pi)
  else
    if (exists_v (dom pi)) and (exists_u (Nu_dom_pi v))
    then routing_tokens_to_destination (no_token_swap v u)
    else routing_tokens_to_destination (unhappy_swap)


Require Export Utilities.
Require Export SQIR.UnitaryOps.

Local Open Scope ucom.

(****************************)
(**   Program Definition   **)
(****************************)

(** Quantum Phase Estimation (QPE) program definition **)

(* Controlled rotation cascade on n qubits. *)
Fixpoint controlled_rotations n : base_ucom n :=
  match n with
  | 0 | 1 => SKIP
  | 2     => control 1 (Rz (2 * PI / 2 ^ n) 0) (* makes 0,1 cases irrelevant *)
  | S n'  => cast (controlled_rotations n') n ;
            control n' (Rz (2 * PI / 2 ^ n) 0)
  end.

(* Quantum Fourier transform on n qubits.
   We use the definition below (with cast and map_qubits) for proof convenience.
   For a more standard functional definition of QFT see Quipper:
   https://www.mathstat.dal.ca/~selinger/quipper/doc/src/Quipper/Libraries/QFT.html *)
Fixpoint QFT n : base_ucom n :=
  match n with
  | 0    => SKIP
  | 1    => H 0
  | S n' => H 0 ; controlled_rotations n ;
           cast (map_qubits S (QFT n')) n
  end.

(* The QFT puts the qubits in the wrong order, so you typically want to reverse
   them at the end. *)
Fixpoint reverse_qubits' dim n : base_ucom dim :=
  match n with
  | 0    => SKIP
  | 1    => SWAP 0 (dim - 1) (* makes 0 case irrelevant *)
  | S n' => reverse_qubits' dim n' ; SWAP n' (dim - n' - 1)
  end.
Definition reverse_qubits n := reverse_qubits' n (n/2)%nat.
Definition QFT_w_reverse n := QFT n ; reverse_qubits n.

Fixpoint controlled_powers' {n} (c : base_ucom n) k kmax : base_ucom (kmax + n) :=
  match k with
  | 0    => SKIP
  | 1    => cast (control (kmax - 1) c) (kmax + n) (* makes 0 case irrelevant *)
  | S k' => controlled_powers' c k' kmax ;
           cast (niter (2 ^ k') (control (kmax - k' - 1) c)) (kmax + n)
  end.
Definition controlled_powers {n} (c : base_ucom n) k := controlled_powers' c k k.

(* k = number of digits in result
   n = number of qubits in input state *)
Definition QPE k n (c : base_ucom n) : base_ucom (k + n) :=
  cast (npar k U_H) (k + n) ;
  controlled_powers (map_qubits (fun q => k + q)%nat c) k;
  cast (invert (QFT_w_reverse k)) (k + n).

Fixpoint controlled_powers_var' {n} (f : nat -> base_ucom n) k kmax : base_ucom (kmax + n) :=
  match k with
  | 0    => SKIP
  | 1    => cast (control (kmax - 1) (f O)) (kmax + n) (* makes 0 case irrelevant *)
  | S k' => controlled_powers_var' f k' kmax ;
           cast (control (kmax - k' - 1) (f k')) (kmax + n)
  end.
Definition controlled_powers_var {n} (f : nat -> base_ucom n) k := controlled_powers_var' f k k.

Definition QPE_var k n (f : nat -> base_ucom n) : base_ucom (k + n) :=
  cast (npar k U_H) (k + n) ;
  controlled_powers_var (fun x => map_qubits (fun q => k + q)%nat (f x)) k;
  cast (invert (QFT_w_reverse k)) (k + n).

(****************************)
(**         Proofs         **)
(****************************)

(** Well-typedness of QFT (used in proof of QPE) **)

Local Opaque control.
Local Transparent Rz.
Lemma controlled_rotations_WT : forall n,
  (n > 1)%nat -> uc_well_typed (controlled_rotations n).
Proof.
  intros n Hn.
  destruct n; try lia.
  destruct n; try lia.
  induction n.
  simpl.
  apply uc_well_typed_control; repeat split; try lia.
  1,2: constructor; auto.
  replace (controlled_rotations (S (S (S n)))) with (cast (controlled_rotations (S (S n))) (S (S (S n))) ; control (S (S n)) (Rz (2 * PI / 2 ^ (S (S (S n)))) 0)) by reflexivity.
  constructor.
  apply typed_cast; try lia.
  apply IHn; try lia.
  apply uc_well_typed_control; repeat split; try lia.
  1,2: constructor; lia.
Qed.
Local Opaque Rz.
Local Transparent control.

Lemma QFT_WT : forall n, (n > 0)%nat -> uc_well_typed (QFT n).
Proof.
  intros n Hn.
  destruct n; try lia.
  induction n.
  simpl. apply uc_well_typed_H; lia.
  replace (QFT (S (S n))) with (H 0 ; controlled_rotations (S (S n)) ; cast (map_qubits S (QFT (S n))) (S (S n))) by reflexivity.
  repeat constructor.
  apply uc_well_typed_H; lia.
  apply controlled_rotations_WT; lia.
  replace S with (fun i => 1 + i)%nat by reflexivity.
  apply uc_well_typed_map_qubits.
  simpl. apply IHn. lia.
Qed.

Lemma reverse_qubits_WT : forall n, (n > 0)%nat -> uc_well_typed (reverse_qubits n).
Proof.
  assert (H: forall n dim, (n > 0)%nat -> (2 * n <= dim)%nat -> uc_well_typed (reverse_qubits' dim n)).
  { intros n dim Hn Hdim.
    destruct n; try lia.
    induction n.
    apply uc_well_typed_SWAP; lia.
    replace (reverse_qubits' dim (S (S n))) with (reverse_qubits' dim (S n) ; SWAP (S n) (dim - (S n) - 1)) by reflexivity.
    constructor.
    apply IHn; lia.
    apply uc_well_typed_SWAP; lia. }
  intros n Hn.
  bdestruct (n =? 1); subst.
  unfold reverse_qubits; simpl.
  apply uc_well_typed_ID; lia.
  apply H.
  apply Nat.div_str_pos. lia.
  apply Nat.mul_div_le. lia.
Qed.

Lemma QFT_w_reverse_WT : forall n, (n > 0)%nat -> uc_well_typed (QFT_w_reverse n).
Proof. intros. constructor. apply QFT_WT; auto. apply reverse_qubits_WT; auto. Qed.

(** Proof of QFT semantics **)

Definition b2R (b : bool) : R := if b then 1%R else 0%R.
Local Coercion b2R : bool >-> R.
Lemma f_to_vec_controlled_Rz : forall (n i j : nat) (θ : R) (f : nat -> bool),
  (i < n)%nat -> (j < n)%nat -> (i <> j)%nat ->
  uc_eval (control i (Rz θ j)) × (f_to_vec n f)
      = (Cexp (f i * f j * θ)) .* f_to_vec n f.
Proof.
  intros.
  rewrite control_correct; auto.
  rewrite Mmult_plus_distr_r.
  rewrite Mmult_assoc.
  rewrite f_to_vec_Rz by auto.
  rewrite Mscale_mult_dist_r.
  destruct (f i) eqn:fi.
  rewrite (f_to_vec_proj_eq _ _ _ true) by auto.
  rewrite (f_to_vec_proj_neq _ _ _ false); auto.
  2: rewrite fi; easy.
  Msimpl_light.
  rewrite Rmult_1_l.
  reflexivity.
  rewrite (f_to_vec_proj_eq _ _ _ false) by auto.
  rewrite (f_to_vec_proj_neq _ _ _ true); auto.
  2: rewrite fi; easy.
  simpl.
  autorewrite with R_db Cexp_db.
  Msimpl_light.
  reflexivity.
  Local Transparent Rz.
  1,2: constructor; auto.
  Local Opaque Rz.
Qed.

Lemma controlled_rotations_action_on_basis : forall n f,
  (n > 1)%nat ->
  (uc_eval (controlled_rotations n)) × (f_to_vec n f) =
    Cexp (2 * PI * (f O) * INR (funbool_to_nat (n-1) (shift f 1%nat)) / (2 ^ n)) .*
      (f_to_vec n f).
Proof.
  intros n f Hn.
  destruct n; try lia.
  destruct n; try lia.
  induction n.
  - simpl uc_eval.
    rewrite f_to_vec_controlled_Rz by lia.
    apply f_equal2; try reflexivity.
    unfold funbool_to_nat, shift; simpl.
    destruct (f (S O)); destruct (f O); simpl;
    autorewrite with R_db; reflexivity.
  - (* easier way to do the following? "simpl" produces a gross expression;
       also, too much manual rewriting below *)
    replace (uc_eval (controlled_rotations (S (S (S n))))) with (uc_eval (control (S (S n)) (Rz (2 * PI / 2 ^ (S (S (S n)))) 0)) × uc_eval (cast (controlled_rotations (S (S n))) (S (S (S n))))) by reflexivity.
    replace (f_to_vec (S (S (S n))) f) with (f_to_vec (S (S n)) f ⊗ ∣ Nat.b2n (f (S (S n))) ⟩) by reflexivity.
    replace (S (S (S n))) with (S (S n) + 1)%nat by lia.
    rewrite <- pad_dims_r.
    simpl I.
    rewrite Mmult_assoc.
    replace (2 ^ ((S (S n)) + 1))%nat with (2 ^ (S (S n)) * (2 ^ 1))%nat by unify_pows_two.
    restore_dims.
    rewrite kron_mixed_product.
    rewrite IHn by lia.
    Msimpl.
    rewrite Mscale_kron_dist_l.
    replace (f_to_vec (S (S n)) f ⊗ ∣ Nat.b2n (f (S (S n))) ⟩) with (f_to_vec (S (S (S n))) f) by reflexivity.
    replace (2 ^ ((S (S n)) + 1))%nat with (2 ^ (S (S n)) * (2 ^ 1))%nat by unify_pows_two.
    rewrite Mscale_mult_dist_r.
    replace (1 * 1)%nat with 1%nat by reflexivity.
    replace (2 ^ (S (S n)) * (2 ^ 1))%nat with (2 ^ ((S (S n)) + 1))%nat by unify_pows_two.
    replace (2 ^ (S (S n)) * 2)%nat with (2 ^ ((S (S n)) + 1))%nat by unify_pows_two.
    replace (S (S (S n))) with (S (S n) + 1)%nat by lia.
    rewrite f_to_vec_controlled_Rz by lia.
    rewrite Mscale_assoc.
    apply f_equal2; try reflexivity.
    rewrite <- Cexp_add.
    replace (f (S (S n)) * f O * (2 * PI / 2 ^ (S (S n) + 1)))%R with (2 * PI * f O * f (S (S n)) * / 2 ^ (S (S n) + 1))%R by lra.
    autorewrite with R_db.
    repeat rewrite Rmult_assoc.
    repeat rewrite <- Rmult_plus_distr_l.
    repeat rewrite <- Rmult_assoc.
    simpl.
    replace (INR (funbool_to_nat (S n) (shift f 1)) * / (2 * (2 * 2 ^ n)) + f (S (S n)) * / (2 * (2 * 2 ^ (n + 1))))%R with (INR (funbool_to_nat (S (n + 1)) (shift f 1)) * / (2 * (2 * 2 ^ (n + 1))))%R.
    repeat rewrite Rmult_assoc.
    reflexivity.
    replace (n + 1)%nat with (S n) by lia.
    unfold funbool_to_nat.
    simpl binlist_to_nat.
    repeat rewrite plus_INR.
    unfold shift; simpl.
    field_simplify_eq; try nonzero.
    replace (S (n + 1)) with (S (S n)) by lia.
    destruct (f (S (S n))); simpl; lra.
    apply controlled_rotations_WT; lia.
Qed.

Lemma QFT_semantics : forall n f,
  (n > 0)%nat ->
  uc_eval (QFT n) × (f_to_vec n f) =
    / √(2 ^ n) .* vkron n (fun i => ∣0⟩ .+ Cexp (2 * PI * INR (funbool_to_nat (n - i) (shift f i)) / 2 ^ (n - i)) .* ∣1⟩).
Proof.
  intros n f Hn.
  generalize dependent f.
  destruct n; try lia.
  induction n; intro f.
  - simpl QFT.
    rewrite f_to_vec_H by lia.
    simpl; Msimpl.
    unfold funbool_to_nat, shift; simpl.
    destruct (f O); simpl; autorewrite with R_db; try reflexivity.
    replace (2 * PI * / 2)%R with PI by lra.
    reflexivity.
  - replace (QFT (S (S n))) with (H 0 ; controlled_rotations (S (S n)) ; cast (map_qubits S (QFT (S n))) (S (S n))) by reflexivity.
    Local Opaque QFT controlled_rotations Nat.pow funbool_to_nat.
    simpl uc_eval.
    repeat rewrite Mmult_assoc.
    rewrite f_to_vec_H by lia.
    distribute_scale. distribute_plus. distribute_scale.
    rewrite 2 controlled_rotations_action_on_basis by lia.
    rewrite 2 update_index_eq.
    rewrite 2 Mscale_mult_dist_r.
    specialize (pad_dims_l (QFT (S n)) (S O)) as H.
    simpl in H.
    replace (fun q : nat => S q) with S in H by reflexivity.
    rewrite <- H; clear H.
    rewrite 2 (f_to_vec_split 0 (S (S n)) 0) by lia.
    remember (S n) as n'.
    simpl f_to_vec.
    replace (0 + 0)%nat with O by reflexivity.
    rewrite 2 update_index_eq.
    Msimpl_light.
    replace (n' - 0 - 0)%nat with n' by lia.
    repeat rewrite kron_mixed_product.
    Msimpl_light.
    rewrite 2 f_to_vec_shift_update_oob by lia.
    assert (H: (n' > 0)%nat) by lia.
    specialize (IHn H (shift f 1)).
    rewrite IHn; clear - Heqn'.
    distribute_scale.
    repeat rewrite <- Mscale_kron_dist_l.
    rewrite <- kron_plus_distr_r.
    simpl Cexp at 1 3; autorewrite with R_db C_db Cexp_db.
    rewrite Cmult_comm.
    rewrite <- Mscale_assoc.
    rewrite <- Mscale_plus_distr_r.
    rewrite <- Mscale_kron_dist_l.
    rewrite Mscale_assoc.
    rewrite Mscale_kron_dist_l.
    apply f_equal2. (* some missing automation from lca/lra here *)
    rewrite <- tech_pow_Rmult.
    rewrite sqrt_mult; try lra.
    rewrite RtoC_mult.
    rewrite Cinv_mult_distr; nonzero.
    apply pow_le; lra.
    replace (n' - 0)%nat with n' by lia.
    rewrite <- Cexp_add.
    replace (shift (update f 0 true) 1) with (shift f 1).
    2: { unfold shift. apply functional_extensionality.
         intro x. rewrite update_index_neq; auto. lia. }
    replace (VectorStates.b2R (f O) * PI + 2 * PI * INR (funbool_to_nat n' (shift f 1)) * / (2 * 2 ^ n'))%R with (2 * PI * INR (funbool_to_nat (S n') f) * / (2 ^ S n'))%R.
    2: { Local Transparent funbool_to_nat.
         rewrite funbool_to_nat_shift with (k:=S O) by lia.
         unfold funbool_to_nat; simpl.
         field_simplify_eq; try nonzero.
         replace (n' - 0)%nat with n' by lia.
         repeat (try rewrite mult_INR; try rewrite plus_INR; try rewrite pow_INR).
         simpl. replace (1 + 1)%R with 2%R by lra.
         destruct (f O); simpl; lra. }
    remember (fun i => ∣0⟩ .+ Cexp (2 * PI * INR (funbool_to_nat (S n' - i) (shift f i)) / 2 ^ (S n' - i)) .* ∣1⟩) as f'.
    replace (fun i => ∣0⟩ .+ Cexp (2 * PI * INR (funbool_to_nat (n' - i) (shift (shift f 1) i)) / 2 ^ (n' - i)) .* ∣1⟩) with (shift f' 1).
    2: { rewrite Heqf'. unfold shift.
         apply functional_extensionality. intro x.
         replace (S n' - (x + 1))%nat with (n' - x)%nat by lia.
         replace (fun i => f (i + (x + 1))%nat) with (fun i => f (i + x + 1)%nat).
         reflexivity.
         apply functional_extensionality. intro x0.
         replace (x0 + x + 1)%nat with (x0 + (x + 1))%nat by lia.
         reflexivity. }
    simpl Nat.b2n.
    replace (∣ 0 ⟩ .+ Cexp (2 * PI * INR (funbool_to_nat (S n') f) * / 2 ^ S n') .* ∣ 1 ⟩) with (f' O).
    2: { subst; simpl. rewrite shift_0. reflexivity. }
    rewrite vkron_extend_l.
    reflexivity.
    intro i; subst; simpl; auto with wf_db.
Qed.
Local Transparent QFT controlled_rotations Nat.pow funbool_to_nat.

(* The property in the previous lemma can be stated without the shift operation *)
Lemma Cexp_shift : forall n i f, (i < n)%nat ->
  Cexp (2 * PI * INR (funbool_to_nat (n - i) (shift f i)) / 2 ^ (n - i)) =
    Cexp (2 * PI * INR (funbool_to_nat n f) / 2 ^ (n - i)).
Proof.
  intros n i f H.
  rewrite (funbool_to_nat_shift n f i) by auto.
  rewrite plus_INR, mult_INR, pow_INR.
  replace (INR 2) with 2%R by (simpl; lra).
  rewrite Rmult_plus_distr_l, Rdiv_plus_distr.
  rewrite Cexp_add.
  replace (2 * PI * (2 ^ (n - i) * INR (funbool_to_nat i f)) / 2 ^ (n - i))%R with (IZR (2 * Z_of_nat (funbool_to_nat i f)) * PI)%R.
  rewrite Cexp_2nPI.
  lca.
  rewrite mult_IZR, <- INR_IZR_INZ.
  field_simplify_eq; nonzero.
Qed.

Lemma SWAP_action_on_product_state :
  forall m n dim (A : Vector (2 ^ m)) (B : Vector (2 ^ (n - m - 1))) (C : Vector (2 ^ (dim - n - 1))) (ψ1 ψ2 : Vector 2),
  (m < n)%nat -> (n < dim)%nat ->
  WF_Matrix ψ1 -> WF_Matrix ψ2 -> WF_Matrix A -> WF_Matrix B -> WF_Matrix C ->
  @ Mmult _ _ (1 * 1 * 1 * 1 * 1) (@uc_eval dim (SWAP m n)) (A ⊗ ψ1 ⊗ B ⊗ ψ2 ⊗ C) = A ⊗ ψ2 ⊗ B ⊗ ψ1 ⊗ C.
Proof.
  intros m n dim A B C ψ1 ψ2 ? ? ? ? ? ? ?.
  autorewrite with eval_db.
  gridify.
  replace ((m + (1 + x + 1) + d1 - (m + 1 + x) - 1))%nat with d1 by lia.
  rewrite (ket_decomposition ψ1) by auto.
  rewrite (ket_decomposition ψ2) by auto.
  autorewrite with ket_db.
  repeat rewrite (Cmult_comm (ψ1 _ _)).
  lma.
Qed.

Lemma SWAP_symmetric : forall m n dim, (@SWAP dim m n) ≡ SWAP n m.
Proof.
  intros. unfold uc_equiv.
  autorewrite with eval_db.
  gridify.
Qed.

Lemma SWAP_on_2_qubits : @uc_eval 2 (SWAP 0 1) = swap.
Proof. autorewrite with eval_db. solve_matrix. Qed.

Lemma SWAP_action_on_vkron : forall dim m n (f : nat -> Vector 2),
  (dim > 0)%nat -> (m < dim)%nat -> (n < dim)%nat -> (m <> n)%nat ->
(*  (forall i, (i < dim)%nat -> WF_Matrix (f i)) -> *)
  (forall i, WF_Matrix (f i)) -> (* Changed *)
  @Mmult _ _ (1 * 1) (uc_eval (SWAP m n)) (vkron dim f) =
    vkron dim (fun k => if (k =? m) then f n else if (k =? n) then f m else f k).
Proof.
  intros dim m n f Hdim Hm Hn Hneq WF.
  destruct dim; try lia.
  destruct dim; try lia.
  destruct dim.
  - destruct n; destruct m; try lia.
    destruct m; try lia. simpl.
    rewrite SWAP_symmetric, SWAP_on_2_qubits.
    Qsimpl. reflexivity.
    destruct n; try lia. simpl.
    rewrite SWAP_on_2_qubits.
    Qsimpl. reflexivity.
  - remember (S (S (S dim))) as dim'.
    remember (fun k : nat => if k =? m then f n else if k =? n then f m else f k) as f'.
    bdestruct (m <? n).
    + assert (WF' : forall i, WF_Matrix (f' i)). subst. intros. destruct (i =? m), (i =? n); auto.
      rewrite 2 (vkron_split dim' m); auto with wf_db; try lia.
      rewrite 2 (vkron_split (dim' - 1 - m) (n - m - 1)); auto with wf_db; try lia.
      replace (dim' - 1 - m - 1 - (n - m - 1))%nat with (dim' - n - 1)%nat by lia. (* slow *)
      restore_dims. repeat rewrite <- kron_assoc by auto with wf_db. restore_dims.
      repeat rewrite shift_simplify.
      rewrite SWAP_action_on_product_state; auto with wf_db.
      repeat rewrite shift_plus.
      replace (n - m - 1 + (m + 1))%nat with n by lia.
      replace (n - m - 1 + 1 + (m + 1))%nat with (n + 1)%nat by lia.
      rewrite (vkron_eq _ f f').
      replace (f n) with (f' m).
      rewrite (vkron_eq _ (shift f (m + 1)) (shift f' (m + 1))).
      replace (f m) with (f' n).
      rewrite (vkron_eq _ (shift f (n + 1)) (shift f' (n + 1))).
      reflexivity.
      all: intros; subst f'; unfold shift; bdestruct_all; trivial.
      all: try apply vkron_WF; intros; apply WF; lia.
    + assert (WF' : forall i, WF_Matrix (f' i)). subst. intros. destruct (i =? m), (i =? n); auto.
      rewrite 2 (vkron_split dim' n); auto with wf_db; try lia.
      rewrite 2 (vkron_split (dim' - 1 - n) (m - n - 1)); auto with wf_db; try lia.
      replace (dim' - 1 - n - 1 - (m - n - 1))%nat with (dim' - m - 1)%nat by lia.
      restore_dims. repeat rewrite <- kron_assoc by auto with wf_db. restore_dims.
      repeat rewrite shift_simplify.
      rewrite SWAP_symmetric.
      rewrite SWAP_action_on_product_state; auto with wf_db; try lia.
      repeat rewrite shift_plus.
      replace (m - n - 1 + (n + 1))%nat with m by lia.
      replace (m - n - 1 + 1 + (n + 1))%nat with (m + 1)%nat by lia.
      rewrite (vkron_eq _ f f').
      replace (f m) with (f' n).
      rewrite (vkron_eq _ (shift f (n + 1)) (shift f' (n + 1))).
      replace (f n) with (f' m).
      rewrite (vkron_eq _ (shift f (m + 1)) (shift f' (m + 1))).
      reflexivity.
      all: intros; subst f'; unfold shift; bdestruct_all; trivial.
      all: try apply vkron_WF; intros; apply WF; lia.
Qed.

Lemma reverse_qubits'_action_on_vkron : forall dim n (f : nat -> Vector 2),
  (n > 0)%nat -> (2 * n <= dim)%nat ->
  (forall i : nat, WF_Matrix (f i)) ->
  @Mmult _ _ 1 (uc_eval (reverse_qubits' dim n)) (vkron dim f) =
    vkron dim (fun k => if ((k <? n)%nat || (dim - n - 1 <? k)%nat)
                     then f (dim - k - 1)%nat else f k).
Proof.
  intros dim n f Hn1 Hn2 WF.
  induction n; try lia.
  destruct n.
  - clear IHn. simpl.
    rewrite SWAP_action_on_vkron; try lia.
    apply vkron_eq.
    intros.
    bdestruct_all; simpl.
    subst. replace (dim - 0 - 1)%nat with (dim - 1)%nat by lia. reflexivity.
    subst. replace (dim - (dim - 1) - 1)%nat with O by lia. reflexivity.
    reflexivity.
    apply WF.
  - replace (uc_eval (reverse_qubits' dim (S (S n)))) with (uc_eval (SWAP (S n) (dim - (S n) - 1)) × uc_eval (reverse_qubits' dim (S n))) by reflexivity.
    rewrite Mmult_assoc.
    rewrite IHn by lia; clear IHn.
    rewrite SWAP_action_on_vkron; try lia.
    apply vkron_eq.
    intros.
    bdestruct_all; simpl; subst; trivial.
    replace (dim - (dim - S n - 1) - 1)%nat with (S n) by lia.
    reflexivity.
    intros.
    bdestruct_all; simpl; apply WF; lia.
Qed.

Lemma reverse_qubits_action_on_vkron : forall n f,
  (n > 1)%nat -> (forall i : nat, WF_Matrix (f i)) ->
  uc_eval (reverse_qubits n) × (vkron n f) = vkron n (fun k => f (n - k - 1)%nat).
Proof.
  intros n f Hn WF.
  unfold reverse_qubits.
  rewrite reverse_qubits'_action_on_vkron.
  apply vkron_eq.
  intros i Hi.
  bdestruct (i <? n / 2); bdestruct (n - n / 2 - 1 <? i); simpl; try reflexivity.
  destruct (Nat.even n) eqn:evn.
  apply Nat.even_spec in evn.
  destruct evn. subst.
  rewrite (Nat.mul_comm 2 x) in *.
  try rewrite Nat.div_mul in *; lia.
  apply negb_true_iff in evn.
  apply Nat.odd_spec in evn.
  destruct evn. subst.
  replace ((2 * x + 1) / 2)%nat with x in *.
  replace (2 * x + 1 - i - 1)%nat with i by lia. reflexivity.
  rewrite Nat.add_comm.
  replace (S O) with (Nat.b2n true) by reflexivity.
  rewrite Nat.add_b2n_double_div2. reflexivity.
  apply Nat.div_str_pos. lia.
  apply Nat.mul_div_le. lia.
  apply WF.
Qed.

(* QFT w/ reverse takes basis state ∣x⟩ to (1/√N) \sum_{k=0}^{N-1} e^{2πixk/N} ∣k⟩;
   this is a more useful form of QFT_semantics *)
Lemma QFT_w_reverse_semantics : forall n (f : nat -> bool),
  (n > 1)%nat ->
  uc_eval (QFT_w_reverse n) × (f_to_vec n f) =
    / √(2 ^ n) .* big_sum (fun k => Cexp (2 * PI * INR (funbool_to_nat n f * k) / (2 ^ n)) .* basis_vector (2^n) k) (2^n).
Proof.
  intros n f Hn.
  unfold QFT_w_reverse; simpl.
  rewrite Mmult_assoc.
  rewrite QFT_semantics by lia.
  distribute_scale.
  rewrite reverse_qubits_action_on_vkron; auto with wf_db.
  apply f_equal2; try reflexivity.
  remember (2 * PI * INR (funbool_to_nat n f) / 2 ^ n)%R as c.
  rewrite (vkron_eq _ _ (fun k : nat => ∣0⟩ .+ Cexp (c * 2 ^ (n - k - 1)) .* ∣1⟩)).
  rewrite (big_sum_eq_bounded _ (fun k : nat => Cexp (c * INR k) .* basis_vector (2 ^ n) k)).
  apply vkron_to_vsum1. lia.
  intros i Hi.
  subst. apply f_equal2; try reflexivity. apply f_equal.
  rewrite mult_INR. lra.
  intros i Hi.
  rewrite Cexp_shift by lia.
  subst. do 2 (apply f_equal2; try reflexivity). apply f_equal.
  field_simplify_eq; try nonzero.
  repeat rewrite Rmult_assoc.
  rewrite <- pow_add.
  replace (n - (n - i - 1) + (n - i - 1))%nat with n by lia.
  reflexivity.
Qed.

Lemma Csum_geometric_series : forall (c : C) (n : nat),
  1 - c <> 0 -> Σ (fun i => c ^ i) n = (1 - c ^ n) / (1 - c).
Proof.
  intros c n Hc.
  induction n; simpl. lca.
  rewrite IHn.
  field_simplify_eq; try lca. (* lca should complete proof here... *)
  apply Hc.
Qed.

Lemma Cexp_neq_1 : forall x,
  x <> 0 -> -2 * PI < x < 2 * PI -> Cexp x <> 1.
Proof.
  intros x Hnz [Hlt Hgt] contra.
  apply pair_equal_spec in contra as [H _].
  assert (cos x < 1).
  { destruct (Rlt_le_dec x 0).
    destruct (Rlt_le_dec x (- PI)).
    rewrite <- cos_2PI, <- cos_neg.
    apply cos_increasing_1; lra.
    rewrite <- cos_0, <- cos_neg.
    apply cos_decreasing_1; lra.
    destruct (Rlt_le_dec x PI).
    rewrite <- cos_0.
    apply cos_decreasing_1; lra.
    rewrite <- cos_2PI.
    apply cos_increasing_1; lra. }
  lra.
Qed.

Lemma Csum_Cexp_nonzero : forall (z : BinInt.Z) (N : nat),
  (IZR z <> 0) -> (- INR N < IZR z < INR N)%R ->
  Σ (fun i => Cexp (2 * PI * IZR z / INR N) ^ i) N = 0.
Proof.
  intros z N Hnz [Hineq1 Hineq2].
  rewrite Csum_geometric_series.
  rewrite Cexp_pow.
  replace (2 * PI * IZR z / INR N * INR N)%R with (IZR (2 * z) * PI)%R.
  rewrite Cexp_2nPI. lca.
  rewrite mult_IZR. field_simplify_eq; try nonzero.
  assert (H : Cexp (2 * PI * IZR z / INR N) <> 1).
  apply Cexp_neq_1.
  (* is there a tactic that can prove inequalities over R? *)
  do 3 (apply Rmult_integral_contrapositive_currified; try nonzero).
  apply PI_neq0.
  split.
  assert (2 * PI * (- INR N) / INR N < 2 * PI * IZR z / INR N).
  replace (2 * PI * - INR N / INR N)%R with (- INR N * (2 * PI / INR N))%R.
  replace (2 * PI * IZR z / INR N)%R with (IZR z * (2 * PI / INR N))%R.
  apply Rmult_lt_compat_r.
  apply Rmult_lt_0_compat.
  apply Rgt_2PI_0.
  nonzero.
  assumption.
  1,2: field_simplify_eq; nonzero.
  replace (2 * PI * - INR N / INR N)%R with (-2 * PI)%R in H.
  2: field_simplify_eq; nonzero.
  assumption.
  replace (2 * PI * IZR z / INR N)%R with (IZR z * (2 * PI / INR N))%R.
  replace (2 * PI)%R with (INR N * (2 * PI / INR N))%R at 2.
  apply Rmult_lt_compat_r.
  apply Rmult_lt_0_compat.
  apply Rgt_2PI_0.
  nonzero.
  assumption.
  1,2: field_simplify_eq; nonzero.
  apply Cminus_eq_contra.
  apply not_eq_sym.
  assumption.
Qed.

Lemma QFT_w_reverse_semantics_inverse : forall n (f : nat -> bool),
  (n > 1)%nat ->
  uc_eval (invert (QFT_w_reverse n)) × f_to_vec n f =
    (/ √(2 ^ n) .* big_sum (fun k => Cexp (- 2 * PI * INR (funbool_to_nat n f * k) / (2 ^ n)) .* basis_vector (2^n) k) (2^n)).
Proof.
  intros n f Hn.
  rewrite <- invert_correct.
  rewrite <- (Mmult_1_l _ _ (_ .* _)); auto with wf_db.
  assert (H : (n > 0)%nat) by lia.
  specialize (uc_eval_unitary n (QFT_w_reverse n) (QFT_w_reverse_WT n H)) as [_ WFU].
  rewrite <- WFU.
  rewrite Mmult_assoc.
  apply f_equal2; try reflexivity.
  distribute_scale. rewrite Mmult_Msum_distr_l.
  erewrite big_sum_eq_bounded.
  2: { intros i Hi.
       distribute_scale.
       rewrite basis_f_to_vec_alt by auto.
       rewrite QFT_w_reverse_semantics by auto.
       rewrite Mscale_assoc.
       rewrite <- Mscale_Msum_distr_r.
       erewrite big_sum_eq_bounded.
       2: { intros i0 Hi0.
            rewrite Mscale_assoc.
            rewrite (Cmult_comm _ (/ _)).
            rewrite <- Cmult_assoc, <- Cexp_add.
            rewrite nat_to_funbool_inverse by auto.
            replace (-2 * PI * INR (funbool_to_nat n f * i) / 2 ^ n + 2 * PI * INR (i * i0) / 2 ^ n)%R with ((2 * PI * (INR i0 - INR (funbool_to_nat n f)) / 2 ^ n) * INR i)%R.
            reflexivity.
            repeat rewrite mult_INR; lra. }
       reflexivity. }
  rewrite big_sum_swap_order.
  rewrite (big_sum_eq_bounded _ (fun i => / √ (2 ^ n) .* (if i =? funbool_to_nat n f then (2 ^ n) .* basis_vector (2 ^ n) i else Zero))).
  2: { intros i Hi.
       rewrite Mscale_Msum_distr_l.
       rewrite <- Csum_mult_l.
       rewrite <- Mscale_assoc.
       apply f_equal2; try reflexivity.
       bdestruct (i =? funbool_to_nat n f).
       subst.
       rewrite Csum_1.
       rewrite pow_INR; simpl. replace (1 + 1)%R with 2%R by lra.
       rewrite <- RtoC_pow. reflexivity.
       intro x.
       simpl. rewrite <- Cexp_0. apply f_equal. lra.
       replace (fun x : nat => Cexp (2 * PI * (INR i - INR (funbool_to_nat n f)) / 2 ^ n * INR x)) with (fun x => Cexp (2 * PI * (IZR (Z.of_nat i - Z.of_nat (funbool_to_nat n f))) / INR (2 ^ n)) ^ x).
       rewrite Csum_Cexp_nonzero. Msimpl. reflexivity.
       rewrite minus_IZR. rewrite <- 2 INR_IZR_INZ.
       apply Rminus_eq_contra.
       apply not_INR. assumption.
       specialize (funbool_to_nat_bound n f) as Hf.
       clear - Hi Hf.
       apply inj_lt in Hi. apply inj_lt in Hf.
       replace (INR (2 ^ n)) with (IZR (Z.of_nat (2 ^ n))).
       split.
       rewrite <- opp_IZR. apply IZR_lt. lia.
       apply IZR_lt. lia.
       rewrite <- INR_IZR_INZ, pow_INR; simpl.
       replace (1 + 1)%R with 2%R by lra. reflexivity.
       apply functional_extensionality; intro x.
       rewrite Cexp_pow.
       rewrite minus_IZR. rewrite <- 2 INR_IZR_INZ.
       rewrite pow_INR.
       reflexivity. }
  rewrite Mscale_Msum_distr_r.
  rewrite Mscale_assoc.
  replace (big_sum (fun i : nat => if i =? funbool_to_nat n f then (2 ^ n) .* basis_vector (2 ^ n) i else Zero) (2 ^ n)) with (2 ^ n .* basis_vector (2 ^ n) (funbool_to_nat n f)).
  rewrite basis_f_to_vec.
  rewrite Mscale_assoc.
  rewrite <- (Mscale_1_l _ _ (basis_vector _ _)) at 1.
  apply f_equal2; try reflexivity.
  field_simplify_eq; try nonzero.
  rewrite Csqrt_sqrt. rewrite RtoC_pow. reflexivity.
  rewrite pow_IZR. apply IZR_le. apply Z.pow_nonneg. lia.
  specialize (funbool_to_nat_bound n f) as ?.
  erewrite big_sum_unique.
  reflexivity.
  exists (funbool_to_nat n f).
  repeat split.
  assumption.
  rewrite Nat.eqb_refl. reflexivity.
  intros. bdestruct_all. reflexivity.
Qed.

(** Proof of QPE semantics **)

Lemma f_to_vec_controlled_U : forall n k (c : base_ucom n) (ψ : Vector (2 ^ n)) (θ : R) i j (f : nat -> bool),
  (k > 0)%nat -> (n > 0)%nat -> (j < k)%nat -> (i > 0)%nat ->
  uc_well_typed c -> WF_Matrix ψ ->
  (uc_eval c) × ψ = Cexp θ .* ψ ->
  @Mmult _ _ (1 * 1) (uc_eval (cast (niter i (control j (map_qubits (fun q : nat => (k + q)%nat) c))) (k + n))) ((f_to_vec k f) ⊗ ψ) =
    Cexp (f j * INR i * θ) .* (f_to_vec k f) ⊗ ψ.
Proof.
  intros n k c ψ θ i j f ? ? ? ? WT WF Heig.
  rewrite cast_niter_commute.
  rewrite niter_correct by lia.
  rewrite cast_control_commute.
  rewrite <- niter_correct by lia.
  rewrite niter_control_commute by lia.
  rewrite control_correct; try lia.
  rewrite Mmult_plus_distr_r.
  rewrite Mmult_assoc.
  rewrite niter_correct by lia.
  rewrite <- pad_dims_l.
  replace (2 ^ (k + n))%nat with (2 ^ k * 2 ^ n)%nat by unify_pows_two.
  rewrite Mmult_n_kron_distr_l.
  rewrite kron_mixed_product.
  rewrite Mmult_n_1_r.
  Msimpl.
  erewrite Mmult_n_eigenvector; auto.
  2: apply Heig.
  distribute_scale.
  replace (proj j (k + n) false) with (proj j k false ⊗ I (2 ^ n)).
  2: unfold proj; autorewrite with eval_db; gridify; trivial.
  replace (proj j (k + n) true) with (proj j k true ⊗ I (2 ^ n)).
  2: unfold proj; autorewrite with eval_db; gridify; trivial.
  restore_dims.
  repeat rewrite kron_mixed_product.
  Msimpl.
  destruct (f j) eqn:fj.
  rewrite (f_to_vec_proj_eq _ _ _ true) by auto.
  rewrite (f_to_vec_proj_neq _ _ _ false); auto.
  2: rewrite fj; easy.
  Msimpl.
  rewrite Rmult_1_l.
  rewrite Cexp_pow.
  rewrite Rmult_comm.
  reflexivity.
  rewrite (f_to_vec_proj_eq _ _ _ false) by auto.
  rewrite (f_to_vec_proj_neq _ _ _ true); auto.
  2: rewrite fj; easy.
  simpl. autorewrite with R_db Cexp_db.
  Msimpl.
  reflexivity.
  apply is_fresh_niter; auto.
  apply map_qubits_fresh; auto.
  apply uc_well_typed_niter.
  apply uc_well_typed_map_qubits; auto.
Qed.

Lemma niter_1 : forall {d} (c : base_ucom d), niter 1 c = c.
Proof. reflexivity. Qed.

Lemma controlled_powers'_action_on_basis :
  forall k kmax n (c : base_ucom n) (ψ : Vector (2^n)) f θ,
  (n > 0)%nat -> (k > 0)%nat -> (kmax >= k)%nat -> uc_well_typed c -> WF_Matrix ψ ->
  (uc_eval c) × ψ = Cexp (2 * PI * θ) .* ψ ->
  @Mmult _ _ (1 * 1) (uc_eval (controlled_powers' (map_qubits (fun q => kmax + q)%nat c) k kmax)) ((f_to_vec kmax f) ⊗ ψ) =
    Cexp (2 * PI * θ * INR (funbool_to_nat k (shift f (kmax - k)))) .* ((f_to_vec kmax f) ⊗ ψ).
Proof.
  intros k kmax n c ψ f θ Hn Hk Hkmax WT WF Heigen.
  destruct k; try lia.
  induction k.
  - simpl.
    rewrite <- (niter_1 (cast _ _)).
    rewrite <- cast_niter_commute.
    erewrite f_to_vec_controlled_U; try apply Heigen; auto.
    2: lia.
    rewrite Mscale_kron_dist_l.
    apply f_equal2; try reflexivity.
    unfold funbool_to_nat, shift; simpl.
    rewrite Nat.add_0_r.
    apply f_equal.
    destruct (f (kmax - 1)%nat); simpl; lra.
  - replace (controlled_powers' (map_qubits (fun q : nat => (kmax + q)%nat) c) (S (S k)) kmax) with (controlled_powers' (map_qubits (fun q : nat => (kmax + q)%nat) c) (S k) kmax ; cast (niter (2 ^ (S k)) (control (kmax - (S k) - 1) (map_qubits (fun q : nat => (kmax + q)%nat) c))) (kmax + n)) by reflexivity.
    Local Opaque controlled_powers' Nat.pow.
    simpl uc_eval.
    rewrite Mmult_assoc.
    rewrite IHk by lia; clear IHk.
    distribute_scale.
    erewrite f_to_vec_controlled_U; try apply Heigen; auto.
    2,3: lia.
    2: assert (2 ^ S k <> 0)%nat by (apply Nat.pow_nonzero; lia); lia.
    restore_dims. distribute_scale.
    apply f_equal2; try reflexivity.
    rewrite <- Cexp_add.
    apply f_equal.
    rewrite (funbool_to_nat_shift (S (S k)) _ (S O)) by lia.
    replace (S (S k) - 1)%nat with (S k) by lia.
    replace (shift (shift f (kmax - S (S k))) 1) with (shift f (kmax - S k)).
    2: { unfold shift. apply functional_extensionality; intro x.
         replace (x + 1 + (kmax - S (S k)))%nat with (x + (kmax - S k))%nat by lia.
         reflexivity. }
    rewrite plus_INR, mult_INR. repeat rewrite pow_INR.
    unfold shift; simpl.
    replace (1 + 1)%R with 2%R by lra.
    replace (kmax - S (S k))%nat with (kmax - S k - 1)%nat by lia.
    unfold funbool_to_nat; simpl.
    field_simplify_eq.
    destruct (f (kmax - S k - 1)%nat); simpl; lra.
Qed.

Lemma controlled_powers_action_on_basis :
  forall k n (c : base_ucom n) (ψ : Vector (2^n)) f θ,
  (n > 0)%nat -> (k > 0)%nat -> uc_well_typed c -> WF_Matrix ψ ->
  (uc_eval c) × ψ = Cexp (2 * PI * θ) .* ψ ->
  @Mmult _ _ (1 * 1) (uc_eval (controlled_powers (map_qubits (fun q => k + q)%nat c) k)) ((f_to_vec k f) ⊗ ψ) =
  Cexp (2 * PI * θ * INR (funbool_to_nat k f)) .* ((f_to_vec k f) ⊗ ψ).
Proof.
  intros k n c ψ f θ Hn Hk WT WF Heigen.
  unfold controlled_powers.
  erewrite controlled_powers'_action_on_basis; try apply Heigen; auto.
  replace (k - k)%nat with O by lia.
  rewrite shift_0.
  reflexivity.
Qed.

(* Simplify the expression uc_eval (QPE k n c) × (k ⨂ ∣0⟩ ⊗ ψ) *)
Local Opaque QFT_w_reverse Nat.mul Nat.pow.
Lemma QPE_simplify : forall k n (c : base_ucom n) (ψ : Vector (2 ^ n)) θ,
  (n > 0)%nat -> (k > 1)%nat -> uc_well_typed c -> WF_Matrix ψ ->
  (uc_eval c) × ψ = Cexp (2 * PI * θ) .* ψ ->
  @Mmult _ _ (1 * 1) (uc_eval (QPE k n c)) (k ⨂ ∣0⟩ ⊗ ψ) =
    (/ (2 ^ k) .* big_sum (fun i : nat => (Σ (fun j => Cexp (2 * PI * (θ - INR i / 2 ^ k) * INR j)) (2 ^ k)) .* basis_vector (2 ^ k) i) (2 ^ k) ⊗ ψ).
Proof.
  intros k n c ψ θ Hn Hk WT WF Heig.
  unfold QPE; simpl.
  repeat rewrite Mmult_assoc.
  repeat rewrite <- pad_dims_r.
  rewrite npar_H by lia.
  replace (2 ^ (k + n))%nat with (2 ^ k * 2 ^ n)%nat by unify_pows_two.
  replace (1 * 1)%nat with (1 ^ k * 1)%nat.
  2: rewrite Nat.pow_1_l; reflexivity.
  rewrite kron_mixed_product.
  Msimpl.
  rewrite H0_kron_n_spec_alt by lia.
  restore_dims. distribute_scale.
  rewrite kron_Msum_distr_r.
  replace (2 ^ (k + n))%nat with (2 ^ k * 2 ^ n)%nat by unify_pows_two.
  rewrite Mmult_Msum_distr_l.
  erewrite big_sum_eq_bounded.
  2: { intros i Hi.
       rewrite basis_f_to_vec_alt by auto. restore_dims.
       erewrite controlled_powers_action_on_basis; try apply Heig; auto; try lia.
       rewrite nat_to_funbool_inverse by auto. reflexivity. }
  rewrite Mmult_Msum_distr_l.
  erewrite big_sum_eq_bounded.
  2: { intros i Hi.
       rewrite Mscale_mult_dist_r.
       rewrite kron_mixed_product.
       Msimpl.
       rewrite QFT_w_reverse_semantics_inverse by auto.
       distribute_scale.
       rewrite (Cmult_comm _ (/ _)).
       rewrite <- Mscale_assoc.
       rewrite <- Mscale_kron_dist_l.
       rewrite <- Mscale_Msum_distr_r.
       erewrite big_sum_eq_bounded.
       2: { intros i0 Hi0.
            rewrite Mscale_assoc.
            rewrite <- Cexp_add.
            rewrite nat_to_funbool_inverse by auto.
            replace (2 * PI * θ * INR i + -2 * PI * INR (i * i0) / 2 ^ k)%R with ((2 * PI * (θ - INR i0 / 2 ^ k)) * INR i)%R.
            reflexivity.
            repeat rewrite mult_INR; lra. }
       rewrite <- Mscale_kron_dist_l.
       reflexivity. }
  rewrite <- kron_Msum_distr_r.
  rewrite Mscale_Msum_distr_r.
  distribute_scale.
  rewrite big_sum_swap_order.
  erewrite big_sum_eq_bounded.
  2: { intros i Hi.
       rewrite Mscale_Msum_distr_l.
       reflexivity. }
  replace (/ √ (2 ^ k) * / √ (2 ^ k)) with (/ 2 ^ k).
  reflexivity.
  rewrite <- Cinv_mult_distr; try nonzero.
  rewrite <- RtoC_mult, sqrt_def, <- RtoC_pow.
  reflexivity.
  apply pow_le. lra.
  apply npar_WT; try lia.
  rewrite <- uc_well_typed_invert.
  apply QFT_w_reverse_WT.
  lia.
Qed.
Local Transparent QFT_w_reverse Nat.mul.

(** Simplified QPE - for the general proof see QPEGeneral.v

  Preconditions:
   - z is the k-bit dyadic rational representation of θ
   - U × ∣ψ⟩ = Cexp (2πθ / 2^k) .* ∣ψ⟩

  Postcondition: the first k bits of the output state are z *)
Lemma QPE_semantics_simplified : forall k n (c : base_ucom n) z (ψ : Vector (2 ^ n)),
  (n > 0)%nat -> (k > 1)%nat -> uc_well_typed c -> WF_Matrix ψ ->
  let θ := (INR (funbool_to_nat k z) / 2 ^ k)%R in
  (uc_eval c) × ψ = Cexp (2 * PI * θ) .* ψ ->
  @Mmult _ _ (1 * 1) (uc_eval (QPE k n c)) (k ⨂ ∣0⟩ ⊗ ψ) = (f_to_vec k z) ⊗ ψ.
Proof.
  intros k n c z ψ Hn Hk WT WF θ Heig.
  rewrite QPE_simplify with (θ := θ) by assumption.
  rewrite (big_sum_eq_bounded _ (fun i => if i =? funbool_to_nat k z then (2 ^ k) .* basis_vector (2 ^ k) i else Zero)).
  specialize (funbool_to_nat_bound k z) as Hz.
  rewrite (big_sum_unique (2 ^ k .* basis_vector (2 ^ k) (funbool_to_nat k z))).
  distribute_scale.
  replace (/ 2 ^ k * 2 ^ k) with C1.
  rewrite Mscale_1_l.
  rewrite basis_f_to_vec. reflexivity.
  rewrite Cinv_l; nonzero.
  exists (funbool_to_nat k z).
  repeat split.
  assumption.
  rewrite Nat.eqb_refl. reflexivity.
  intros. bdestruct_all. reflexivity.
  intros i Hi.
  bdestruct (i =? funbool_to_nat k z).
  rewrite Csum_1.
  rewrite RtoC_pow, pow_INR; simpl. reflexivity.
  intro x.
  subst i θ.
  simpl. rewrite <- Cexp_0.
  apply f_equal. lra.
  replace (fun j => Cexp (2 * PI * (θ - INR i / 2 ^ k) * INR j)) with (fun j => Cexp (2 * PI * IZR (Z.of_nat (funbool_to_nat k z) - Z.of_nat i) / INR (2 ^ k)) ^ j).
  rewrite Csum_Cexp_nonzero.
  Msimpl. reflexivity.
  rewrite minus_IZR, <- 2 INR_IZR_INZ.
  apply Rminus_eq_contra.
  apply not_INR. apply not_eq_sym. assumption.
  specialize (funbool_to_nat_bound k z) as Hz.
  clear - Hi Hz.
  apply inj_lt in Hi. apply inj_lt in Hz.
  replace (INR (2 ^ k))%R with (IZR (Z.of_nat (2 ^ k))).
  split.
  rewrite <- opp_IZR. apply IZR_lt. lia.
  apply IZR_lt. lia.
  rewrite <- INR_IZR_INZ, pow_INR; simpl.
  replace (1 + 1)%R with 2%R by lra. reflexivity.
  apply functional_extensionality; intro x.
  rewrite Cexp_pow.
  rewrite minus_IZR, <- 2 INR_IZR_INZ.
  subst θ. apply f_equal.
  rewrite pow_INR.
  replace (INR 2) with 2 by reflexivity.
  lra.
Qed.

(* QPE_var is the variant of QPE used in our implementation of Shor's.
   It takes a circuit constructor (f : nat -> base_ucom), which has the
   same effect as (fun i => niter (2^i) c) when acting on eigenstates.
*)

Local Transparent Nat.pow controlled_powers'.

Lemma f_to_vec_controlled_U_var : forall n k (fc : nat -> base_ucom n) (ψ : Vector (2 ^ n)) (θ : R) i j (f : nat -> bool),
  (k > 0)%nat -> (n > 0)%nat -> (j < k)%nat -> (i > 0)%nat ->
  uc_well_typed (fc i) -> WF_Matrix ψ ->
  (uc_eval (fc i)) × ψ = Cexp (θ * (INR (2^i))) .* ψ ->
  @Mmult _ _ (1 * 1) (uc_eval (cast (control j (map_qubits (fun q : nat => (k + q)%nat) (fc i))) (k + n)%nat)) ((f_to_vec k f) ⊗ ψ) =
    Cexp (f j * θ * INR (2^i)) .* (f_to_vec k f) ⊗ ψ.
Proof.
  intros n k c ψ θ i j f ? ? ? ? WT WF Heig.
  rewrite cast_control_commute.
  rewrite control_correct; try lia.
  rewrite Mmult_plus_distr_r.
  rewrite Mmult_assoc.
  rewrite <- pad_dims_l.
  replace (2 ^ (k + n))%nat with (2 ^ k * 2 ^ n)%nat by unify_pows_two.
  rewrite kron_mixed_product.
  Msimpl.
  distribute_scale.
  replace (proj j (k + n) false) with (proj j k false ⊗ I (2 ^ n)).
  2: unfold proj; autorewrite with eval_db; gridify; trivial.
  replace (proj j (k + n) true) with (proj j k true ⊗ I (2 ^ n)).
  2: unfold proj; autorewrite with eval_db; gridify; trivial.
  restore_dims.
  repeat rewrite kron_mixed_product.
  Msimpl.
  destruct (f j) eqn:fj.
  rewrite (f_to_vec_proj_eq _ _ _ true) by auto.
  rewrite (f_to_vec_proj_neq _ _ _ false); auto.
  2: rewrite fj; easy.
  Msimpl.
  rewrite Rmult_1_l.
  rewrite Heig.
  distribute_scale.
  reflexivity.
  rewrite (f_to_vec_proj_eq _ _ _ false) by auto.
  rewrite (f_to_vec_proj_neq _ _ _ true); auto.
  2: rewrite fj; easy.
  simpl. autorewrite with R_db Cexp_db.
  Msimpl.
  reflexivity.
  apply map_qubits_fresh; auto.
  apply uc_well_typed_map_qubits; auto.
Qed.

Lemma controlled_powers_var'_action_on_basis :
  forall k kmax n (fc : nat -> base_ucom n) (ψ : Vector (2^n)) f θ,
    (n > 0)%nat -> (k > 0)%nat -> (kmax >= k)%nat ->
    (forall i, (i < k)%nat -> uc_well_typed (fc i)) ->
    WF_Matrix ψ ->
    (forall i, (i < k)%nat -> (uc_eval (fc i)) × ψ = Cexp (2 * PI * θ * (INR (2^i))) .* ψ) ->
    @Mmult _ _ (1 * 1) (uc_eval (controlled_powers_var' (fun i => map_qubits (fun q => kmax + q)%nat (fc i)) k kmax)) ((f_to_vec kmax f) ⊗ ψ) =
    Cexp (2 * PI * θ * INR (funbool_to_nat k (shift f (kmax - k)))) .* ((f_to_vec kmax f) ⊗ ψ).
Proof.
  intros k kmax n fc ψ f θ Hn Hk Hkmax WT WF Heigen.
  destruct k; try lia.
  induction k.
  - simpl.
    rewrite <- (niter_1 (cast _ _)).
    rewrite <- cast_niter_commute.
    erewrite f_to_vec_controlled_U; try apply Heigen; auto.
    2: lia.
    rewrite Mscale_kron_dist_l.
    apply f_equal2; try reflexivity.
    unfold funbool_to_nat, shift; simpl.
    rewrite Nat.add_0_r.
    apply f_equal.
    destruct (f (kmax - 1)%nat); simpl; lra.
  - remember (S k) as k'.
    simpl. subst.
    Local Opaque controlled_powers_var' Nat.pow.
    simpl uc_eval.
    rewrite Mmult_assoc.
    restore_dims.
    assert (forall i : nat, (i < S k)%nat -> uc_well_typed (fc i)) by (intros; apply WT; lia).
    assert (forall i : nat, (i < S k)%nat -> uc_eval (fc i) × ψ = Cexp (2 * PI * θ * INR (2 ^ i)) .* ψ) by (intros; apply Heigen; lia).
    rewrite IHk by (try easy; lia). clear IHk.
    distribute_scale.
    erewrite f_to_vec_controlled_U_var; try apply Heigen; auto.
    2,3,4: lia.
    restore_dims. distribute_scale.
    apply f_equal2; try reflexivity.
    rewrite <- Cexp_add.
    apply f_equal.
    rewrite (funbool_to_nat_shift (S (S k)) _ (S O)) by lia.
    replace (S (S k) - 1)%nat with (S k) by lia.
    replace (shift (shift f (kmax - S (S k))) 1) with (shift f (kmax - S k)).
    2: { unfold shift. apply functional_extensionality; intro x.
         replace (x + 1 + (kmax - S (S k)))%nat with (x + (kmax - S k))%nat by lia.
         reflexivity. }
    rewrite plus_INR, mult_INR. repeat rewrite pow_INR.
    unfold shift; simpl.
    replace (1 + 1)%R with 2%R by lra.
    replace (kmax - S (S k))%nat with (kmax - S k - 1)%nat by lia.
    unfold funbool_to_nat; simpl.
    field_simplify_eq.
    destruct (f (kmax - S k - 1)%nat); simpl; lra.
Qed.

Lemma controlled_powers_var_action_on_basis :
  forall k n (fc : nat -> base_ucom n) (ψ : Vector (2^n)) f θ,
    (n > 0)%nat -> (k > 0)%nat ->
    (forall i, (i < k)%nat -> uc_well_typed (fc i)) ->
    WF_Matrix ψ ->
    (forall i, (i < k)%nat -> (uc_eval (fc i)) × ψ = Cexp (2 * PI * θ * (INR (2^i))) .* ψ) ->
    @Mmult _ _ (1 * 1) (uc_eval (controlled_powers_var (fun i => map_qubits (fun q => k + q)%nat (fc i)) k)) ((f_to_vec k f) ⊗ ψ) =
    Cexp (2 * PI * θ * INR (funbool_to_nat k f)) .* ((f_to_vec k f) ⊗ ψ).
Proof.
  intros k n c ψ f θ Hn Hk WT WF Heigen.
  unfold controlled_powers_var.
  erewrite controlled_powers_var'_action_on_basis; try apply Heigen; auto.
  replace (k - k)%nat with O by lia.
  rewrite shift_0.
  reflexivity.
Qed.

Lemma QPE_var_equivalent :
  forall k n (c : base_ucom n) (f : nat -> base_ucom n) (ψ : Vector (2^n)) θ,
    (n > 0)%nat ->
    WF_Matrix ψ ->
    uc_well_typed c ->
    (uc_eval c) × ψ = Cexp (2 * PI * θ) .* ψ ->
    (forall i, (i < k)%nat -> uc_well_typed (f i)) ->
    (forall i, (i < k)%nat -> (uc_eval (f i)) × ψ = Cexp (2 * PI * θ * (INR (2^i))) .* ψ) ->
    @Mmult _ _ (1 * 1) (uc_eval (QPE k n c)) (k ⨂ qubit0 ⊗ ψ) = @Mmult _ _ (1 * 1) (uc_eval (QPE_var k n f)) (k ⨂ qubit0 ⊗ ψ).
Proof.
  assert (G: forall {n} (A B : base_ucom n), uc_eval (A; B) = uc_eval B × uc_eval A) by intuition.
  intros k n c0 f0.
  intros.
  unfold QPE, QPE_var.
  repeat rewrite G.
  repeat rewrite Mmult_assoc.
  bdestruct (k <=? 0). assert (k = O) by lia. subst. easy.
  rewrite <- pad_dims_r with (c := npar k U_H) by (apply npar_WT; lia).
  rewrite npar_H by lia.
  replace (2 ^ (k + n))%nat with (2 ^ k * 2 ^ n)%nat by unify_pows_two.
  replace (1 * 1)%nat with (1 ^ k * 1)%nat.
  2: rewrite Nat.pow_1_l; reflexivity.
  rewrite kron_mixed_product.
  Msimpl.
  rewrite H0_kron_n_spec_alt by lia.
  restore_dims. distribute_scale.
  rewrite kron_Msum_distr_r.
  replace (2 ^ (k + n))%nat with (2 ^ k * 2 ^ n)%nat by unify_pows_two.
  rewrite Mmult_Msum_distr_l.
  rewrite Mmult_Msum_distr_l with (f := (fun i : nat => basis_vector (2 ^ k) i ⊗ ψ)).
  do 2 apply f_equal.
  apply big_sum_eq_bounded. intros.
  rewrite basis_f_to_vec_alt by easy.
  restore_dims.
  rewrite controlled_powers_action_on_basis with (θ := θ) by (try easy; try lia).
  rewrite controlled_powers_var_action_on_basis with (θ := θ); try easy; try lia.
Qed.

Local Transparent controlled_powers_var' SKIP ID.
Lemma controlled_powers_var_WT : forall n (f : nat -> base_ucom n) k,
  (k > 0)%nat ->
  (forall i, uc_well_typed (f i)) ->
  uc_well_typed (controlled_powers_var
                  (fun x : nat => map_qubits (fun q : nat => (k + q)%nat) (f x)) k).
Proof.
  intros n f k Hk WT.
  unfold controlled_powers_var.
  assert (forall kmax, (kmax >= k)%nat ->
    uc_well_typed (controlled_powers_var'
                    (fun x : nat => map_qubits (fun q : nat => (kmax + q)%nat) (f x)) k kmax)).
  { intros kmax Hkmax.
    induction k. lia.
    simpl. destruct k.
    rewrite cast_control_commute.
    apply uc_well_typed_control.
    split; [| split].
    lia.
    apply map_qubits_fresh. lia.
    apply uc_well_typed_map_qubits.
    apply WT.
    constructor.
    apply IHk. lia. lia.
    rewrite cast_control_commute.
    apply uc_well_typed_control.
    split; [| split].
    lia.
    apply map_qubits_fresh. lia.
    apply uc_well_typed_map_qubits.
    apply WT. }
  apply H. lia.
Qed.

Lemma QPE_var_WT : forall k n (f : nat -> base_ucom n),
  (k > 0)%nat ->
  (forall i, uc_well_typed (f i)) ->
  uc_well_typed (QPE_var k n f).
Proof.
  intros k n f Hk WT.
  unfold QPE_var.
  constructor. constructor.
  apply typed_cast.
  apply npar_WT.
  auto. lia.
  apply controlled_powers_var_WT; auto.
  apply typed_cast.
  rewrite <- uc_well_typed_invert.
  apply QFT_w_reverse_WT.
  auto. lia.
Qed.


Require Import ZArith.

Require Import Statement ImprovedPredicative Wpr.

Require Import Lia.

Open Scope stmt_scope.

Definition spec := ⟨fun '(i,n) '(i',n') => i <= n /\ i' = n /\ n = n'⟩.

Definition prog := WWhile (fun '(i,n) => i <> n) Do $(i,n) := (i+1,n) Done.

Theorem correctness : prog ⊑ spec.
Proof.
  intros (i,n) ((u,v),(HHin,_)); clear u v.
  set (K := wpr prog (pred spec)).
  generalize i HHin; clear i HHin.
  induction n; intros i HHin.
  { apply wpr_while_construct; right; simpl. lia. }
  { apply Lt.le_lt_or_eq in HHin.
    destruct HHin as [ HHin | HHin ].
    { assert (i <= n) as HHin' by lia.
      cut ((fun '(i,n) '(i',n') => i' <= n' /\ K (i,S n) (i',S n')) (i, n) (i, n)).
      { intros (HH1,HH2); auto. }
      { apply (IHn _ HHin'); clear i n HHin HHin' IHn.
        intros (i,n) (i',n'). intros [ (HHin,(HHi'n',HHind)) | HH ]; split; try lia.
        { apply wpr_while_construct; left; split; auto. lia. }
        { simpl in HH; apply wpr_while_construct; left; split; try lia; fold prog K.
          apply wpr_while_construct; right; simpl. lia.
        }
      }
    }
    { apply wpr_while_construct; right; simpl. lia. }
  }
Qed.


Inductive natural : Type := Succ : natural -> natural |  Zero : natural.

Inductive lst : Type := Cons : natural -> lst -> lst |  Nil : lst.

Inductive tree : Type := Node : natural -> tree -> tree -> tree |  Leaf : tree.

Inductive Pair : Type := mkpair : natural -> natural -> Pair
with Zlst : Type := zcons : Pair -> Zlst -> Zlst |  znil : Zlst.

Fixpoint plus (plus_arg0 : natural) (plus_arg1 : natural) : natural
           := match plus_arg0, plus_arg1 with
              | Zero, n => n
              | Succ n, m => Succ (plus n m)
              end.
(* Requires No helper lemma *)
Theorem theorem0 : forall (x : natural) (y : natural), eq (plus x (Succ (Succ y))) (Succ (Succ (plus x y))).
Proof.
   intros.
  induction x.
  - simpl. rewrite IHx. reflexivity.
  - reflexivity.
Qed.


(** Generated by coq-of-ocaml *)
Require Import OCaml.OCaml.

Local Set Primitive Projections.
Local Open Scope string_scope.
Local Open Scope Z_scope.
Local Open Scope type_scope.
Import ListNotations.

Unset Positivity Checking.
Unset Guard Checking.

Inductive nat : Set :=
| O : nat
| S : nat -> nat.

Inductive natural : Set :=
| Zero : natural
| Succ : natural -> natural.

Fixpoint plus (plus_arg0 : natural) (plus_arg1 : natural) {struct plus_arg0}
  : natural :=
  match plus_arg0 with
  | Zero => plus_arg1
  | Succ n => Succ (plus n plus_arg1)
  end.

Definition synth (lf1 : natural) (y : natural) : natural :=
  plus Zero (plus y lf1).


Require Export TopologicalSpaces.
Require Export Subbases.
Require Export Continuity.
Require Export Nets.
From ZornsLemma Require Export InverseImage.
From ZornsLemma Require Import FiniteIntersections.

(* Also called "initial topology". Its construction is dual
   (in the categorical sense) to the construction of the strong topology. *)

Section WeakTopology.

Variable X:Type.
Variable A:Type.
Variable Y:A->TopologicalSpace.
Variable f:forall a:A, X->point_set (Y a).

Inductive weak_topology_subbasis : Family X :=
  | intro_fa_inv_image: forall (a:A) (V:Ensemble (point_set (Y a))),
    open V -> In weak_topology_subbasis (inverse_image (f a) V).

Definition WeakTopology : TopologicalSpace :=
  Build_TopologicalSpace_from_subbasis X weak_topology_subbasis.

Lemma weak_topology_makes_continuous_funcs:
  forall a:A, continuous (f a) (X:=WeakTopology).
Proof.
intro.
red; intros.
pose proof (Build_TopologicalSpace_from_subbasis_subbasis
  _ weak_topology_subbasis).
apply H0.
constructor; trivial.
Qed.

Lemma weak_topology_is_weakest: forall (T':Family X)
  (H1:_) (H2:_) (H3:_),
  (forall a:A, continuous (f a)
     (X := Build_TopologicalSpace X T' H1 H2 H3)) ->
  forall U:Ensemble X, @open WeakTopology U -> T' U.
Proof.
intros.
destruct H0.
apply H1.
intros.
apply H0 in H4.
induction H4.
- exact H3.
- destruct H4.
  apply H; trivial.
- apply H2; trivial.
Qed.

Lemma weak_topology_continuous_char (W : TopologicalSpace)
      (g : (point_set W) -> (point_set (WeakTopology))) :
      continuous g <-> (forall a, continuous (compose (f a) g)).
Proof.
split.
- intros.
  unfold compose. apply continuous_composition.
  + apply weak_topology_makes_continuous_funcs.
  + assumption.
- intros.
  apply continuous_subbasis with weak_topology_subbasis.
  { apply Build_TopologicalSpace_from_subbasis_subbasis. }
  intros. induction H0.
  rewrite <- inverse_image_composition.
  apply H. assumption.
Qed.

Section WeakTopology_and_Nets.

Variable I:DirectedSet.
Hypothesis I_nonempty: inhabited (DS_set I).
Variable x:Net I WeakTopology.
Variable x0:X.

Lemma net_limit_in_weak_topology_impl_net_limit_in_projections :
  net_limit x x0 ->
  forall a:A, net_limit (fun i:DS_set I => (f a) (x i)) ((f a) x0).
Proof.
intros.
apply continuous_func_preserves_net_limits; trivial.
apply continuous_func_continuous_everywhere.
apply weak_topology_makes_continuous_funcs.
Qed.

Lemma net_limit_in_projections_impl_net_limit_in_weak_topology :
  (forall a:A, net_limit (fun i:DS_set I => (f a) (x i))
                         ((f a) x0)) ->
  net_limit x x0.
Proof.
intros.
red; intros.
assert (@open_basis WeakTopology
        (finite_intersections weak_topology_subbasis)).
{ apply Build_TopologicalSpace_from_open_basis_basis. }
destruct (open_basis_cover _ H2 x0 U) as [V [? [? ?]]]; trivial.
assert (for large i:DS_set I, In V (x i)).
{ clear H4.
  induction H3.
  - destruct I_nonempty.
    exists X0; constructor.
  - destruct H3.
    destruct H5.
    apply eventually_impl_base with (fun i:DS_set I => In V (f a (x i))).
    + intros.
      constructor; trivial.
    + apply H; trivial.
  - apply eventually_impl_base with
        (fun i:DS_set I => In U0 (x i) /\ In V (x i)).
    + intros.
      destruct H6.
      constructor; trivial.
    + destruct H5.
      apply eventually_and;
        (apply IHfinite_intersections || apply IHfinite_intersections0);
        trivial.
}
refine (eventually_impl_base _ _ _ H6).
intro; apply H4.
Qed.

End WeakTopology_and_Nets.

End WeakTopology.

Arguments WeakTopology {X} {A} {Y}.
Arguments weak_topology_subbasis {X} {A} {Y}.

Require Import ClassicalChoice.

Section WeakTopology1.

Variable X:Type.
Variable Y:TopologicalSpace.
Variable f:X->point_set Y.

Definition WeakTopology1 := WeakTopology (True_rect f).

Lemma weak_topology1_makes_continuous_func:
  continuous f (X:=WeakTopology1).
Proof.
exact (weak_topology_makes_continuous_funcs _ _ _ (True_rect f) I).
Qed.

Lemma weak_topology1_topology:
  forall U:Ensemble X, @open WeakTopology1 U <->
  exists V:Ensemble (point_set Y), open V /\ U = inverse_image f V.
Proof.
split.
2: {
  intros.
  destruct H as [V []].
  subst.
  apply weak_topology1_makes_continuous_func.
  assumption.
}
intros.
red in H.
simpl in H.
destruct H.
assert (forall U:Ensemble X,
  In (finite_intersections (weak_topology_subbasis (True_rect f))) U ->
  exists V:Ensemble (point_set Y), open V /\ U = inverse_image f V).
{ intros.
  induction H0.
  - exists Full_set.
    split.
    + apply open_full.
    + symmetry. apply inverse_image_full.
  - destruct H0.
    destruct a.
    simpl.
    exists V.
    split; trivial.
  - destruct IHfinite_intersections as [V1 [? ?]].
    destruct IHfinite_intersections0 as [V2 [? ?]].
    exists (Intersection V1 V2).
    split.
    + auto with topology.
    + rewrite H3; rewrite H5.
      rewrite inverse_image_intersection.
      trivial.
}
destruct (choice (fun (U:{U:Ensemble X | In F U}) (V:Ensemble (point_set Y))
  => open V /\ proj1_sig U = inverse_image f V)) as [choice_fun].
{ intros.
  destruct x as [U].
  simpl.
  apply H0; auto with sets.
}
exists (IndexedUnion choice_fun).
split.
{ apply open_indexed_union.
  apply H1.
}
apply Extensionality_Ensembles; split; red; intros.
- destruct H2.
  constructor.
  exists (exist _ S H2).
  pose proof (H1 (exist _ S H2)).
  destruct H4.
  simpl in H5.
  rewrite H5 in H3.
  destruct H3.
  exact H3.
- destruct H2.
  inversion H2.
  pose proof (H1 a).
  destruct H5.
  destruct a as [U].
  exists U; trivial.
  simpl in H6.
  rewrite H6.
  constructor.
  exact H3.
Qed.

Lemma weak_topology1_topology_closed:
  forall U:Ensemble X, @closed WeakTopology1 U <->
  exists V:Ensemble (point_set Y), closed V /\ U = inverse_image f V.
Proof.
unfold closed.
split.
- intros. unfold closed in *.
  rewrite weak_topology1_topology in H.
  destruct H as [V []].
  exists (Complement V).
  rewrite inverse_image_complement.
  rewrite <- H0.
  rewrite !Complement_Complement.
  auto.
- intros. destruct H as [V []].
  subst. apply continuous_closed; try assumption.
  apply weak_topology1_makes_continuous_func.
Qed.

Lemma weak_topology1_continuous_char (Z : TopologicalSpace)
      (g : Z -> WeakTopology1) :
  continuous g <->
  continuous (compose f g).
Proof.
  replace f with (True_rect f I).
  2: { reflexivity. }
  unfold WeakTopology1.
  rewrite weak_topology_continuous_char.
  split; intros; auto.
  destruct a. assumption.
Qed.

End WeakTopology1.

Arguments WeakTopology1 {X} {Y}.


Require Export FunctionProperties.
Require Import DecidableDec.
Require Import Description.
Require Import Classical.

(* CSB = Cantor-Schroeder-Bernstein theorem *)

Section CSB.
Variable X Y:Type.
Variable f:X->Y.
Variable g:Y->X.
Hypothesis f_inj: injective f.
Hypothesis g_inj: injective g.

Inductive X_even: X->Prop :=
  | not_g_img: forall x:X, (forall y:Y, g y <> x) -> X_even x
  | g_Y_odd: forall y:Y, Y_odd y -> X_even (g y)
with Y_odd: Y->Prop :=
  | f_X_even: forall x:X, X_even x -> Y_odd (f x).
Inductive X_odd: X->Prop :=
  | g_Y_even: forall y:Y, Y_even y -> X_odd (g y)
with Y_even: Y->Prop :=
  | not_f_img: forall y:Y, (forall x:X, f x <> y) -> Y_even y
  | f_X_odd: forall x:X, X_odd x -> Y_even (f x).

Scheme X_even_coind := Minimality for X_even Sort Prop
  with Y_odd_coind := Minimality for Y_odd Sort Prop.
Scheme X_odd_coind := Minimality for X_odd Sort Prop
  with Y_even_coind := Minimality for Y_even Sort Prop.

Lemma even_odd_excl: forall x:X, ~(X_even x /\ X_odd x).
Proof.
intro.
assert (X_even x -> ~ X_odd x).
2:tauto.
pose proof (X_even_coind (fun x:X => ~ X_odd x) (fun y:Y => ~ Y_even y)).
apply H.
intuition.
destruct H1.
apply H0 with y.
reflexivity.
intuition.
inversion H2.
apply g_inj in H3.
apply H1.
rewrite <- H3.
assumption.
intuition.
inversion H2.
apply H3 with x0.
reflexivity.
apply f_inj in H3.
apply H1.
rewrite <- H3.
assumption.
Qed.

Lemma even_odd_excl2: forall y:Y, ~(Y_even y /\ Y_odd y).
Proof.
intro.
assert (Y_odd y -> ~ Y_even y).
2:tauto.
pose proof (Y_odd_coind (fun x:X => ~ X_odd x) (fun y:Y => ~ Y_even y)).
apply H.
intuition.
destruct H1.
apply H0 with y0.
reflexivity.
intuition.
inversion H2.
apply g_inj in H3.
apply H1.
rewrite <- H3.
assumption.
intuition.
inversion H2.
apply H3 with x.
reflexivity.
apply f_inj in H3.
apply H1.
rewrite <- H3.
assumption.
Qed.

Definition finv: forall y:Y, (exists x:X, f x = y) ->
  { x:X | f x = y }.
intros.
apply constructive_definite_description.
destruct H.
exists x.
red; split.
assumption.
intros.
apply f_inj.
transitivity y; trivial.
symmetry; trivial.
Defined.

Definition ginv: forall x:X, (exists y:Y, g y = x) ->
  { y:Y | g y = x }.
intros.
apply constructive_definite_description.
destruct H.
exists x0.
red; split.
assumption.
intros.
apply g_inj.
transitivity x; trivial; symmetry; trivial.
Defined.

Definition ginv_odd: forall x:X, X_odd x ->
  { y:Y | g y = x }.
intros.
apply ginv.
destruct H.
exists y.
reflexivity.
Defined.

Definition finv_noteven: forall y:Y, ~ Y_even y ->
  { x:X | f x = y }.
intros.
apply finv.
apply NNPP.
red; intro.
contradict H.
constructor 1.
intro; red; intro.
apply H0.
exists x.
assumption.
Defined.

Definition CSB_bijection (x:X) : Y :=
  match (classic_dec (X_odd x)) with
  | left o => proj1_sig (ginv_odd x o)
  | right _ => f x
  end.
Definition CSB_bijection2 (y:Y) : X :=
  match (classic_dec (Y_even y)) with
  | left _ => g y
  | right ne => proj1_sig (finv_noteven y ne)
  end.

Lemma CSB_comp1: forall x:X, CSB_bijection2 (CSB_bijection x) = x.
Proof.
intro.
unfold CSB_bijection; case (classic_dec (X_odd x)).
intro.
destruct ginv_odd.
simpl.
unfold CSB_bijection2; case (classic_dec (Y_even x1)).
intro.
assumption.
intro.
destruct x0.
contradict n.
apply g_inj in e.
rewrite e.
assumption.
intro.
unfold CSB_bijection2; case (classic_dec (Y_even (f x))).
intro.
contradict n.
inversion y.
pose proof (H x).
contradict H1; reflexivity.
apply f_inj in H.
rewrite <- H.
assumption.
intro.
destruct finv_noteven.
simpl.
apply f_inj.
assumption.
Qed.

Lemma CSB_comp2: forall y:Y, CSB_bijection (CSB_bijection2 y) = y.
Proof.
intro.
unfold CSB_bijection2; case (classic_dec (Y_even y)).
intro.
unfold CSB_bijection; case (classic_dec (X_odd (g y))).
intro.
destruct ginv_odd.
simpl.
apply g_inj.
assumption.
intro.
contradict n.
constructor.
assumption.
intro.
destruct finv_noteven.
simpl.
unfold CSB_bijection; case (classic_dec (X_odd x)).
intro.
contradict n.
rewrite <- e.
constructor 2.
assumption.
trivial.
Qed.

Theorem CSB: exists h:X->Y, bijective h.
Proof.
exists CSB_bijection.
apply invertible_impl_bijective.
exists CSB_bijection2.
exact CSB_comp1.
exact CSB_comp2.
Qed.

End CSB.


Require Import Coq.Arith.PeanoNat.
Require Import Coq.Classes.RelationClasses.
Require Export Coq.Lists.List.
Require Import Coq.Sets.Multiset.
Require Export Coq.Sorting.Permutation.
Require Import Coq.Sorting.PermutSetoid.
Require Import Coq.Sorting.PermutEq.
Require Import Omega.

Opaque Nat.eq_dec.

(** Refer to README.md for documentation. *)

(** (foreverbell): generalize to all decidable instances. *)

Lemma list_contents_app_multiplicity_plus :
  forall (a : nat) (l m : list nat),
    multiplicity (list_contents eq Nat.eq_dec (l ++ m)) a =
    multiplicity (list_contents eq Nat.eq_dec l) a +
    multiplicity (list_contents eq Nat.eq_dec m) a.
Proof.
  intros; specialize (list_contents_app eq Nat.eq_dec l m).
  intros. unfold meq, munion in H.
  specialize (H a); auto.
Qed.

Ltac permutation_simplify a (* variable for functional extensionality *) :=
  repeat
    match goal with
    | [ H : Permutation _ _ |- _ ] =>
        rewrite (permutation_Permutation Nat.eq_dec) in H;
        unfold permutation, meq in H;
        specialize (H a);
        repeat (
          simpl list_contents in H;
          unfold munion in H;
          simpl multiplicity in H;
          try rewrite list_contents_app_multiplicity_plus in H
        )
    | [ |- Permutation _ _ ] =>
        rewrite (permutation_Permutation Nat.eq_dec);
        unfold permutation, meq;
        intros a;
        repeat (
          simpl list_contents;
          unfold munion;
          simpl multiplicity;
          try rewrite list_contents_app_multiplicity_plus
        )
    end.

Ltac permutation_solver :=
  let a := fresh "a" in permutation_simplify a;
  repeat
    match goal with
    | [ |- context [if ?A then _ else _] ] => destruct A
    end; omega.


(*common header begin*)
Require Import Utf8.
From Coq Require Import ssreflect ssrfun ssrbool.
Set Implicit Arguments.
Unset Strict Implicit.
Unset Printing Implicit Defensive.
Set Maximal Implicit Insertion.

Require Import List.
Import ListNotations.
(*common header end*)

Require Import PeanoNat.
Require Import Omega.

Require Import ListFacts.
Require Import UserTactics.

(*HintDb containing lemmas regarding the last element of seq*)
Create HintDb seq.

Lemma seq_last : forall (start length : nat),
  seq start (S length) = (seq start length) ++ [start + length].
Proof.
move => start length.
elim : length start.
intros; cbn; f_equal; omega.

move => length IH start; cbn.
f_equal.
have : start + S length = (S start) + length by omega.
move => ->.
apply : IH.
Qed.

Lemma repeat_last : forall (T : Type) (a : T) (length : nat),
  repeat a (S length) = (repeat a length) ++ [a].
Proof.
move => T a length.
elim : length => //.

move => length IH; cbn.
f_equal.
auto.
Qed.

Lemma map_singleton (T U : Type) : forall (f : T -> U) (a : T), map f [a] = [f a].
Proof. reflexivity. Qed.

Lemma Forall_app_singleton : forall (T : Type) (P : T -> Prop) (l : list T) (a : T),
  Forall P (l ++ [a]) <-> Forall P l /\ (P a).
Proof.
intros *.
split.
move /Forall_app.
intuition.
gimme Forall; by inversion.

rewrite ? Forall_forall.
intuition.
gimme In; rewrite ? in_app_iff.
intuition.
gimme In; inversion; done.
Qed.

Hint Rewrite seq_last : seq.
Hint Rewrite @repeat_last : seq.
Hint Rewrite map_app : seq.
Hint Rewrite @map_singleton : seq.
Hint Rewrite plus_O_n plus_n_O : seq.
Hint Rewrite @Forall_app_singleton : seq.
Hint Rewrite <- plus_n_O : seq.

(*list of pars of list elements along with their index starting with start*)
Fixpoint indexed (T : Type) (start : nat) (l : list T) :=
  match l with
  | nil => nil
  | (cons a l) => cons (start, a) (indexed (S start) l)
  end.

Lemma indexed_app : forall (v w : list nat) (n : nat), indexed n (v ++ w) = (indexed n v) ++ (indexed (length v+n) w).
Proof.
elim; cbn; first done.
move => a v IH w n.
have : S (length v + n) = length v + (S n) by omega.
move => ->.
f_equal; eauto.
Qed.

Hint Rewrite @indexed_app : seq.

Lemma in_indexed_in : forall (T : Type) (l : list T) (i n : nat) (a : T), In (i, a) (indexed n l) -> In a l.
Proof.
move => T.
elim; cbn; first done.
move => b l IH i n a.
case; first case; eauto.
Qed.

Lemma in_indexed_bounds : forall (T : Type) (l : list T) (i n : nat) (a : T), In (i, a) (indexed n l) -> n <= i /\ i < n + length l.
Proof.
move => T.
elim; cbn; first done.
move => b l IH i n a.
case; first case.
intros; omega.
move /IH.
case; intros; omega.
Qed.

Lemma in_indexed_eq : forall (T : Type) (l : list T) (i n : nat) (a b : T), In (i, a) (indexed n l) -> In (i, b) (indexed n l) -> a = b.
Proof.
move => T.
elim; cbn; first done.
move => c l IH i n a b.
case.
case=> ? ?.
case.
case=> ? ?; congruence.
subst.
move /in_indexed_bounds; intros; omega.
move => ?.
case.
case => ? ?; subst.
gimme In.
move /in_indexed_bounds; intros; omega.
eauto.
Qed.

Lemma in_in_indexed : forall (T : Type)  (a : T) (l : list T) (n : nat), In a l -> exists (i : nat), In (i, a) (indexed n l).
Proof.
move => T a.
elim => //.
move => b l IH n; cbn.
case.
intro; subst; exists n; by left.
move /IH; move /(_ (S n)). firstorder.
Qed.


(* Copyright (c) 2014, Robert Dockins *)

Require Import basics.
Require Import preord.
Require Import categories.
Require Import sets.
Require Import finsets.
Require Import esets.

Require Import NArith.

(**  * Effective preorders.

       We define the notion of "effective" preorders: those
       for which the order relation is decidable and the
       members of the preorder are enumerable.
  *)

Record effective_order (A:preord) :=
  EffectiveOrder
  { eff_ord_dec : forall x y:A, {x ≤ y} + {x ≰ y}
  ; eff_enum : eset A
  ; eff_complete : forall x:A, x ∈ eff_enum
  }.

(**  Effective orders have a decidable equality.
  *)
Canonical Structure eff_to_ord_dec A (Heff:effective_order A) : ord_dec A :=
  OrdDec A (eff_ord_dec A Heff).
Coercion eff_to_ord_dec : effective_order >-> ord_dec.

(**  The positive integers form an effective preorder.
  *)
Program Definition Ndisc_ord : preord :=
  Preord.Pack N (Preord.Mixin N (@eq N) _ _).
Solve Obligations of Ndisc_ord using intros; subst; auto.
Canonical Structure Ndisc_ord.

Program Definition effective_Nord : effective_order Ndisc_ord
  := EffectiveOrder Ndisc_ord _ (fun n => Some n) _.
Next Obligation.
  simpl. unfold Preord.ord_op. simpl.
  apply N_eq_dec.
Qed.
Next Obligation.
  intros. exists x. auto.
Qed.

(**  Given an effective preorder, we can calculate a canonical
     index value for any given element of the preorder.

     We do this by first generating the subset of the enumeration set
     containing all elements equal to [x]; then we choose the smallest
     index from that set that is defined.  Such an element exists
     because we know the set is inhabited.  The weak principle of countable
     choice then suffices to choose the index.
  *)

Definition unenumerate_set (A:preord) (Heff:effective_order A) (x:A)
  : eset A :=
  fun n =>
    match eff_enum A Heff n with
    | Some x' =>
        if PREORD_EQ_DEC A (eff_to_ord_dec A Heff) x x'
           then Some x'
           else None
    | None => None
    end.

Lemma unenumerate_set_inhabited A Heff x :
  einhabited (unenumerate_set A Heff x).
Proof.
  apply member_inhabited.
  generalize (eff_complete A Heff x).
  intros [n ?].
  case_eq (eff_enum A Heff n); intros.
  rewrite H0 in H.
  exists c. exists n.
  unfold unenumerate_set. rewrite H0.
  destruct (PREORD_EQ_DEC A (eff_to_ord_dec A Heff) x c); auto.
  rewrite H0 in H. elim H.
Qed.

Definition unenumerate (A:preord) (Heff:effective_order A) (x:A) : N :=
  proj1_sig (projT2
    (find_inhabitant A
      (unenumerate_set A Heff x)
      (unenumerate_set_inhabited A Heff x))).

Lemma unenumerate_correct A Heff x :
  exists x', eff_enum A Heff (unenumerate A Heff x) = Some x' /\ x ≈ x'.
Proof.
  unfold unenumerate.
  destruct (find_inhabitant A
               (unenumerate_set A Heff x)
               (unenumerate_set_inhabited A Heff x)); simpl.
  destruct s as [n ?].
  destruct a. simpl.
  case_eq (eff_enum A Heff n); intros.
  unfold unenumerate_set in e.
  rewrite H in e.
  destruct ((PREORD_EQ_DEC A (eff_to_ord_dec A Heff)) x c).
  inversion e. subst x0.
  exists c. split; auto.
  discriminate.
  unfold unenumerate_set in e.
  rewrite H in e. discriminate.
Qed.

(**  The unenumeration index is unique (up to Leibniz equality)
     and equal indexes imply equal elements.
  *)
Lemma unenumerate_uniq A Heff x x' :
  x ≈ x' ->
  unenumerate A Heff x = unenumerate A Heff x'.
Proof.
  intros.
  unfold unenumerate.
  destruct (find_inhabitant A
               (unenumerate_set A Heff x)
               (unenumerate_set_inhabited A Heff x)); simpl.
  destruct (find_inhabitant A
               (unenumerate_set A Heff x')
               (unenumerate_set_inhabited A Heff x')); simpl.
  destruct s as [n [??]].
  destruct s0 as [n' [??]].
  simpl.
  unfold unenumerate_set in e.
  case_eq (eff_enum A Heff n); intros.
  rewrite H0 in e.
  destruct (PREORD_EQ_DEC A (eff_to_ord_dec A Heff) x c).
  inversion e; subst x0. clear e.
  unfold unenumerate_set in e0.
  case_eq (eff_enum A Heff n'); intros.
  rewrite H1 in e0.
  destruct (PREORD_EQ_DEC A (eff_to_ord_dec A Heff) x' c0).
  inversion e0; subst x1. clear e0.
  apply N.le_antisymm.
  apply l with c0; auto.
  unfold unenumerate_set.
  rewrite H1.
  destruct (PREORD_EQ_DEC A (eff_to_ord_dec A Heff) x c0). auto.
  elim n0. rewrite H; auto.
  apply l0 with c; auto.
  unfold unenumerate_set.
  rewrite H0.
  destruct (PREORD_EQ_DEC A (eff_to_ord_dec A Heff) x' c). auto.
  elim n0. rewrite <- H; auto.
  discriminate.
  rewrite H1 in e0. discriminate.
  discriminate.
  rewrite H0 in e. discriminate.
Qed.

Lemma unenumerate_reflects A Heff x x' :
  unenumerate A Heff x = unenumerate A Heff x' ->
  x ≈ x'.
Proof.
  intros.
  unfold unenumerate in *.
  destruct (find_inhabitant A
               (unenumerate_set A Heff x)
               (unenumerate_set_inhabited A Heff x)); simpl in *.
  destruct (find_inhabitant A
               (unenumerate_set A Heff x')
               (unenumerate_set_inhabited A Heff x')); simpl in *.
  destruct s as [n [??]].
  destruct s0 as [m [??]].
  simpl in *.
  subst m.
  unfold unenumerate_set in e.
  unfold unenumerate_set in e0.
  case_eq (eff_enum A Heff n); intros.
  rewrite H in e, e0.
  destruct (PREORD_EQ_DEC A (eff_to_ord_dec A Heff) x c).
  inversion e; subst x0.
  destruct (PREORD_EQ_DEC A (eff_to_ord_dec A Heff) x' c).
  inversion e0; subst x1.
  rewrite e1; auto.
  discriminate.
  discriminate.
  rewrite H in e. discriminate.
Qed.

Lemma eff_in_dec : forall {A:preord} (Heff:effective_order A) (M:finset A) (x:A),
  { x ∈ M } + { x ∉ M }.
Proof.
  intros. apply finset_in_dec.
  apply eff_to_ord_dec. auto.
Qed.

(** The one-point preorder is effective. *)
Program Definition effective_unit : effective_order unitpo
   := EffectiveOrder unitpo (fun _ _ => left I) (single tt) _.
Next Obligation.
  intro. apply single_axiom. destruct x; auto.
Qed.

(** The empty preorder is effective. *)
Program Definition effective_empty : effective_order emptypo :=
  EffectiveOrder _ _ (fun x => None) _.
Next Obligation.
  intros. elim x.
Qed.
Next Obligation.
  intros. elim x.
Qed.

(** The binary product of effective preorders is effective. *)
Program Definition effective_prod {A B:preord}
  (HA:effective_order A)
  (HB:effective_order B)
  : effective_order (A×B)
  := EffectiveOrder _ _ (eprod (eff_enum A HA) (eff_enum B HB)) _.
Next Obligation.
  intros. destruct x; destruct y.
  destruct (eff_ord_dec A HA c c1).
  destruct (eff_ord_dec B HB c0 c2).
  left. split; auto.
  right. intros [??]; apply n; auto.
  right. intros [??]; apply n; auto.
Qed.
Next Obligation.
  intros.
  destruct x as [a b].
  apply eprod_elem.
  split; apply eff_complete.
Qed.

(** The binary sum of effective preorders is effective. *)
Program Definition effective_sum {A B:preord}
  (HA:effective_order A)
  (HB:effective_order B)
  : effective_order (sum_preord A B)
  := EffectiveOrder _ _ (esum (eff_enum A HA) (eff_enum B HB)) _.
Next Obligation.
  intros.
  destruct x; destruct y.
  destruct (eff_ord_dec A HA c c0).
  left. auto.
  right. auto.
  right. intro. elim H.
  right. intro. elim H.
  destruct (eff_ord_dec B HB c c0).
  left. auto.
  right. auto.
Qed.
Next Obligation.
  intros.
  destruct x.
  apply esum_left_elem. apply eff_complete.
  apply esum_right_elem. apply eff_complete.
Qed.

(** The lift of an effective preorder is effective. *)
Definition enum_lift (A:preord) (X:eset A) : eset (lift A) :=
  union2 (single None) (image (liftup A) X).

Program Definition effective_lift {A:preord}
  (HA:effective_order A)
  : effective_order (lift A) :=
  EffectiveOrder _ _ (enum_lift A (eff_enum A HA)) _.
Next Obligation.
  intros.
  destruct x; destruct y; simpl; auto.
  destruct (eff_ord_dec A HA c c0); auto.
  left. hnf. auto.
Qed.
Next Obligation.
  intros. unfold enum_lift.
  apply union2_elem.
  destruct x. right.
  apply image_axiom1'. exists c. split; auto.
  apply eff_complete.
  left. apply single_axiom; auto.
Qed.

(** FIXME? this doesn't really fit here, but the current
    module tree means it can't go in esets.v.  Maybe semidec
    should get split out into a separate file?
  *)
Lemma semidec_ex (A B:preord) (P:A -> B -> Prop)
  (Hok : forall a b c, b ≈ c -> P a b -> P a c)
  (HB:effective_order B) :
  (forall ab, semidec (P (fst ab) (snd ab))) ->
  (forall a, semidec (@ex B (P a))).
Proof.
  intros.
  apply Semidec with (fun n =>
    let (p,q) := pairing.unpairing n in
       match eff_enum B HB p with
       | None => None
       | Some b => decset _ (X (a,b)) q
       end).
  split; simpl; intros.
  destruct H as [n ?].
  case_eq (pairing.unpairing n); intros.
  rewrite H0 in H.
  destruct (eff_enum B HB n0); intros.
  case_eq (decset (P a c) (X (a,c)) n1); intros.
  rewrite H1 in H.
  exists c.
  rewrite <- (decset_correct _ (X (a,c))). simpl.
  hnf; simpl. exists n1. rewrite H1. auto.
  rewrite H1 in H. elim H. elim H.

  destruct H.
  generalize (eff_complete B HB x).
  intros [p ?].
  case_eq (eff_enum B HB p); intros.
  rewrite H1 in H0.
  assert (P a c).
  apply Hok with x; auto.
  rewrite <- (decset_correct _ (X (a,c))) in H2.
  destruct H2 as [q ?].
  simpl in *.
  exists (pairing.pairing (p,q)).
  rewrite pairing.unpairing_pairing.
  rewrite H1.
  destruct (decset (P a c) (X (a,c)) q); auto.
  rewrite H1 in H0. elim H0.
Qed.


Require Export Functor.

Arguments Compose_Functors {_} {_} {_} _ _.
Arguments fmap {_} {_} _ {_} {_} _.
Arguments fobj {_} {_} _ _.

Check fobj.

Class NaturalTransformation (C D: Category)
                            (F  : Functor C D)
                            (G  : Functor C D): Type :=
  mk_nt
  {
    trans    : forall (a: @obj C), (@arrow D (fobj G a) (fobj F a));
    comm_diag: forall {a b: @obj C} (f: arrow  b a), fmap G f o trans a  = trans b o fmap F f;
  }.
Check NaturalTransformation.

Arguments NaturalTransformation {_} {_} _ _.
Arguments trans {_} {_} {_} {_} _ _.

Definition IdNt (C D: Category)
                (F  : Functor C D): NaturalTransformation F F.
Proof. unshelve econstructor.
       - intros. exact (fmap _(@identity C a)).
       - intros. destruct F. simpl. rewrite !preserve_id.
         now rewrite identity_f, f_identity.
Defined.

Program Definition Compose_NaturalTransformations_H
                          (C D E: Category)
                          (F    : Functor C D)
                          (G    : Functor C D)
                          (H    : Functor D E)
                          (I    : Functor D E)
                          (nt1  : NaturalTransformation F G)
                          (nt2  : NaturalTransformation H I):
                                `(NaturalTransformation (Compose_Functors F H)
                                                        (Compose_Functors G I)).
Proof. unshelve econstructor.
       - destruct F, G, H, I, nt1, nt2. simpl in *.
         intros. exact (((fmap2 _ _ (trans0 a)) o trans1 (fobj  a))).
       - destruct F, G, H, I, nt1, nt2. simpl in *.
         intros. rewrite !comm_diag1. rewrite assoc.
         rewrite comm_diag1. rewrite <- assoc.
         rewrite <- !preserve_comp1.
         rewrite <- assoc.
         rewrite <- !preserve_comp1. now rewrite comm_diag0.
Defined.

Program Definition Compose_NaturalTransformations
                      (C D: Category)
                      (F  : @Functor C D)
                      (G  : @Functor C D)
                      (H  : @Functor C D)
                      (nt1: NaturalTransformation F G)
                      (nt2: NaturalTransformation G H): `(NaturalTransformation F H) :=
{|
    trans := fun a: @obj C =>  trans nt2 a o trans nt1 a;
|}.
Next Obligation.
      rewrite assoc.
      rewrite (@comm_diag C D G H nt2 a).
      do 2 rewrite <- assoc.
      rewrite (@comm_diag C D F G nt1 a).
      reflexivity.
Defined.

Arguments Compose_NaturalTransformations {_} {_} {_} {_} {_} _ _.

Lemma Nt_split: forall (C D: Category)
                       (F  : @Functor C D)
                       (G  : @Functor C D)
                       (nt1: NaturalTransformation F G)
                       (nt2: NaturalTransformation F G), trans nt1 = trans nt2 <-> nt1 = nt2.
Proof. intros. split. intros. destruct nt1, nt2, F, G.
       simpl in *. revert comm_diag0. rewrite H. intros.
       specialize (proof_irrelevance (forall (a b : @obj C) (f : arrow b a),
             fmap0 a b f o trans1 a = trans1 b o fmap a b f) comm_diag0 comm_diag1).
       now destruct (proof_irrelevance _ comm_diag0 comm_diag1).
       intros. rewrite H. easy.
Qed.

Class Cone (C D: Category) (F: Functor C D): Type := mkCone
  {
     cobj: @obj D;
     cobl: NaturalTransformation (ConstantFunctor C D cobj) F
  }.

Class Limit (C D: Category) (F: Functor C D): Type := mkLimit
 {
    limc : Cone C D F;
    limob: forall (a b: @obj C) (Cn: Cone C D F), exists !(u: arrow (@cobj C D F limc) (@cobj C D F Cn)),
            trans (@cobl C D F limc) a o u = trans (@cobl C D F Cn) a /\
            trans (@cobl C D F limc) b o u = trans (@cobl C D F Cn) b
 }.

Definition has_Limits (D: Category) := forall (C: Category) (F: Functor C D), Limit C D F.

Definition FunctorCategory (C D: Category): Category.
(* Proof. refine(@mk_Category (@Functor C D F)
                           NaturalTransformation
                           (IdNt C D F)
                           (Compose_NaturalTransformations)
                            _ _ _ ).
       intros. unfold Compose_NaturalTransformations. simpl.
       destruct a, b, c, d, f, g, h. simpl. f_equal.
*)
Proof.
       unshelve econstructor.
       - exact (Functor C D).
       - intros F G. exact (NaturalTransformation G F).
       - intros. exact (@IdNt C D a).
       - intros F G H nt1 nt2.
         exact (Compose_NaturalTransformations nt2 nt1).
       - repeat intro. now rewrite H, H0.
       - intros. apply Nt_split. simpl.
         destruct f, g, h. simpl.
         extensionality a0. now rewrite assoc.
       - intros. apply Nt_split. simpl.
         destruct f, a, b. simpl.
         extensionality a0.
         now rewrite preserve_id0, identity_f.
       - intros. apply Nt_split. simpl.
         destruct f, a, b. simpl.
         extensionality a0.
         now rewrite preserve_id, f_identity.
Defined.

Definition Cat: Category.
Proof. unshelve econstructor.
       - exact Category.
       - intros C D. exact (Functor C D).
       - intro C. exact (@Id C).
       - intros E D C F G. exact (Compose_Functors F G).
       - repeat intro. now subst.
       - intros D C B A F G H. exact (FunctorCompositionAssoc F G H).
       - intros D C F. exact (ComposeIdl F).
       - intros D C F. exact (ComposeIdr F).
Defined.

Class IsomorphismFunctorial {C D: Category} : Type := {
  toC   : Functor C D;
  fromC : Functor D C;

  iso_to_fromC : Compose_Functors toC fromC = @Id C;
  iso_from_toD : Compose_Functors fromC toC = @Id D
}.

Class IsomorphismNT {C D: Category} (F G: Functor C D): Type :=
  mk_IsomorphisnNT
  {
     nt1        : NaturalTransformation F G;
     nt2        : NaturalTransformation G F;
     equivnt_ob1: Compose_NaturalTransformations nt1 nt2 = IdNt C D F;
     equivnt_ob2: Compose_NaturalTransformations nt2 nt1 = IdNt C D G
  }.

Lemma eqIso1: forall (C D: @obj Cat), @Isomorphism Cat C D -> @IsomorphismFunctorial C D.
Proof. intros C D I; destruct I; cbn in *.
       unshelve econstructor.
       - exact iso_from_to.
       - exact iso_to_from.
Qed.

Lemma eqIso2: forall (C D: @obj Cat), @IsomorphismFunctorial C D -> @Isomorphism Cat C D.
Proof. intros C D I.
       unshelve econstructor; destruct I; cbn in *.
       - exact fromC0.
       - exact toC0.
       - exact iso_from_toD0.
       - exact iso_to_fromC0.
Qed.

Lemma eqIso3: forall C D (F G: Functor C D),
                         @Isomorphism (FunctorCategory C D) F G ->
                         @IsomorphismNT C D F G.
Proof. intros C D F G E.
       destruct E; cbn in *.
       - unshelve econstructor.
         + exact iso_from_to.
         + exact iso_to_from.
Qed.

Lemma eqIso4: forall C D (F G: Functor C D),
                         @IsomorphismNT C D F G ->
                         @Isomorphism (FunctorCategory C D) F G.
Proof. intros C D F G I.
       destruct I; unshelve econstructor; cbn in *.
       - exact nt3.
       - exact nt4.
       - exact equivnt_ob4.
       - exact equivnt_ob3.
Qed.

Class EquivalenceCat {C D: Category} (F: Functor C D) (G: Functor D C): Type :=
  mk_EquivalenceCat
  {
     equiv_ob1: @Isomorphism (FunctorCategory C C) (Compose_Functors F G) (@Id C);
     equiv_ob2: @Isomorphism (FunctorCategory D D) (Compose_Functors G F) (@Id D)
  }.

(*
Definition CurryingFunctor (C D E: Category) (F: Functor (Product_Category C D) E):
  Functor C (FunctorCategory D E).
Proof. intros.
       unshelve econstructor.
       - cbn. intro a.
         + unshelve econstructor.
           ++ cbn. destruct F. cbn in *. intro b.
              exact (fobj (a, b)).
           ++ cbn. intros. destruct F. cbn in *.
              clear fmapP preserve_id preserve_comp.
              specialize (fmap (a, a0) (a, b)). cbn in *.
              apply fmap. exact (identity a, f).
           ++ repeat intro. now subst.
           ++ cbn. intros. destruct F. cbn in .
              now rewrite preserve_id.
           ++ cbn. intros. destruct F. cbn in *.
              specialize (preserve_comp (a, a0) (a, b) (a, c)). cbn in *.
              specialize (preserve_comp  (identity a, g) (identity a, f)). cbn in *.
              rewrite identity_f in preserve_comp. now rewrite preserve_comp.
       - intros. cbn.
         unshelve econstructor.
         + intros. cbn. destruct F. cbn in *.
           clear fmapP preserve_id preserve_comp.
           apply fmap. cbn. exact (f, identity a0).
         + intros. cbn. destruct F. cbn in *.
*)

Definition muT(C D   : Category)
              (F     : @Functor C D)
              (G     : @Functor D C)
              (T     := (Compose_Functors F G))
              (T2    := (Compose_Functors T T))
              (epstr : (NaturalTransformation T (@Id C))): (NaturalTransformation T2 T).
Proof. destruct epstr, F, G. simpl in *. unfold T in *.
       refine (@mk_nt C
                      C
                      T2
                      T
                      (fun a => fmap0 _ _ (fmap _ _ (trans0 a))) _).
       intros. unfold T, T2, id in *. simpl in *.
       now rewrite <- !preserve_comp0, <- !preserve_comp, comm_diag0.
Defined.

Definition delD (C D  : Category)
                (F    : @Functor C D)
                (G    : @Functor D C)
                (cT   := (Compose_Functors G F))
                (cT2  := (Compose_Functors cT cT))
                (etatr: (NaturalTransformation (@Id D) cT)): (NaturalTransformation cT cT2).
Proof. destruct etatr, F, G. simpl in *. unfold cT in *.
       refine (@mk_nt D
                      D
                      cT
                      cT2
                      (fun a => fmap _ _ (fmap0 _ _ (trans0 a))) _).
       intros. unfold cT, cT2, id in *. simpl in *.
       now rewrite <- !preserve_comp, <- !preserve_comp0, comm_diag0.
Defined.


Require Export D.

(** **** Problem #1 : 2 stars, optional (poly_exercises) *)
(** Here are a few simple exercises, just like ones in the [Lists]
    chapter, for practice with polymorphism.  Fill in the definitions
    and complete the proofs below. *)

Fixpoint repeat {X : Type} (n : X) (count : nat) : list X :=
  match count with
  | O => nil
  | S m => n :: repeat n m
  end.

Example test_repeat1:
  repeat true 2 = cons true (cons true nil).
Proof. reflexivity. Qed.

Theorem nil_app : forall X:Type, forall l:list X,
  app [] l = l.
Proof.
  reflexivity.
Qed.

Theorem snoc_with_append : forall X : Type,
                         forall l1 l2 : list X,
                         forall v : X,
  snoc (l1 ++ l2) v = l1 ++ (snoc l2 v).
Proof.
  intros. induction l1. reflexivity.
  simpl. rewrite -> IHl1. reflexivity.
Qed.