Require Import GateCancellation.
Require Import HadamardReduction.
Require Import NotPropagation.
Require Import Optimize1qGates.
Require Import RotationMerging.
Require Import RzQGateSet.
Require Import SwapRoute.
Require Import MappingValidation.
Require Import GreedyLayout.
Require Import FullGateSet.
Import FullList.

Local Close Scope Q_scope.
Local Close Scope C_scope.
Local Close Scope R_scope.

(** This file contains the VOQC transformations that are extracted to OCaml,
   along with their correctness properties. The definitions and proofs in this
   file are largely wrappers around definitions and proofs in other files. **)

Definition circ := full_ucom_l.

Definition layout := Layouts.layout.

Definition c_graph : Type := nat * (nat -> nat -> bool).
Definition graph_dim (cg : c_graph) := fst cg.
Definition is_in_graph (cg : c_graph) := snd cg.

Definition path_finding_fun : Type := nat -> nat -> list nat.
Definition qubit_ordering_fun : Type := option nat -> list nat.

(* Cast function changes the dependent type; it will be extracted to a no-op *)
Fixpoint cast {dim} (c : circ dim) dim' : @circ dim' :=
  match c with
  | [] => []
  | App1 g m :: t => App1 g m :: cast t dim'
  | App2 g m n :: t => App2 g m n :: cast t dim'
  | App3 g m n p :: t => App3 g m n p :: cast t dim'
  end.

(** * Utility functions **)

Definition check_well_typed {dim} (c : circ dim) (n : nat) :=
  uc_well_typed_l_b n (cast c n).
Definition convert_to_ibm {dim} (c : circ dim) :=
  FullGateSet.convert_to_ibm c.
Definition convert_to_rzq {dim} (c : circ dim) :=
  FullGateSet.convert_to_rzq c.
Definition replace_rzq {dim} (c : circ dim) :=
  FullGateSet.replace_rzq c.
Definition decompose_to_cnot {dim} (c : circ dim) :=
  FullGateSet.decompose_to_cnot c.

Lemma check_well_typed_correct : forall {dim} (c : circ dim) n,
  check_well_typed c n = true <-> uc_well_typed_l (cast c n).
Proof. intros. apply uc_well_typed_l_b_equiv. Qed.

Lemma convert_to_ibm_preserves_semantics : forall {dim} (c : circ dim),
  (convert_to_ibm c =l= c)%ucom.
Proof. intros. apply FullGateSet.convert_to_ibm_sound. Qed.

Ltac show_preserves_WT H :=
  eapply uc_equiv_l_implies_WT;
  [ symmetry; apply H | assumption ].

Ltac show_preserves_WT_cong H :=
  eapply uc_cong_l_implies_WT;
  [ symmetry; apply H | assumption ].

Lemma convert_to_ibm_preserves_WT : forall {dim} (c : circ dim),
  uc_well_typed_l c -> uc_well_typed_l (convert_to_ibm c).
Proof. intros dim c H. show_preserves_WT (convert_to_ibm_preserves_semantics c). Qed.

Lemma convert_to_ibm_preserves_mapping : forall {dim} (l : full_ucom_l dim) (cg : c_graph),
  respects_constraints_directed (is_in_graph cg) U_CX l ->
  respects_constraints_directed (is_in_graph cg) U_CX (convert_to_ibm l).
Proof.
  intros.
  apply FullGateSet.convert_to_ibm_preserves_mapping.
  assumption.
Qed.

Lemma convert_to_ibm_uses_ibm_gates : forall {dim} (c : circ dim),
  forall_gates only_ibm (convert_to_ibm c).
Proof. intros. apply FullGateSet.convert_to_ibm_gates. Qed.

Lemma convert_to_rzq_preserves_semantics : forall {dim} (c : circ dim),
  (convert_to_rzq c ≅l≅ c)%ucom.
Proof. intros. apply FullGateSet.convert_to_rzq_sound. Qed.

Lemma convert_to_rzq_preserves_WT : forall {dim} (c : circ dim),
  uc_well_typed_l c -> uc_well_typed_l (convert_to_rzq c).
Proof.
  intros dim c H.
  show_preserves_WT_cong (convert_to_rzq_preserves_semantics c).
Qed.

Lemma convert_to_rzq_preserves_mapping : forall {dim} (l : full_ucom_l dim) (cg : c_graph),
  respects_constraints_directed (is_in_graph cg) U_CX l ->
  respects_constraints_directed (is_in_graph cg) U_CX (convert_to_rzq l).
Proof.
  intros.
  apply FullGateSet.convert_to_rzq_preserves_mapping.
  assumption.
Qed.

Lemma convert_to_rzq_uses_rzq_gates : forall {dim} (c : circ dim),
  forall_gates only_rzq (convert_to_rzq c).
Proof. intros. apply FullGateSet.convert_to_rzq_gates. Qed.

Lemma replace_rzq_preserves_semantics : forall {dim} (c : circ dim),
  (replace_rzq c =l= c)%ucom.
Proof. intros. apply FullGateSet.replace_rzq_sound. Qed.

Lemma replace_rzq_preserves_WT : forall {dim} (c : circ dim),
  uc_well_typed_l c -> uc_well_typed_l (replace_rzq c).
Proof. intros dim c H. show_preserves_WT (replace_rzq_preserves_semantics c). Qed.

Lemma replace_rzq_preserves_mapping : forall {dim} (l : full_ucom_l dim) (cg : c_graph),
  respects_constraints_directed (is_in_graph cg) U_CX l ->
  respects_constraints_directed (is_in_graph cg) U_CX (replace_rzq l).
Proof.
  intros.
  apply FullGateSet.replace_rzq_preserves_mapping.
  assumption.
Qed.

Lemma replace_rzq_does_not_use_rzq_gates : forall {dim} (c : circ dim),
  forall_gates no_rzq (replace_rzq c).
Proof. intros. apply FullGateSet.replace_rzq_gates. Qed.

Lemma decompose_to_cnot_preserves_semantics : forall {dim} (c : circ dim),
  (decompose_to_cnot c =l= c)%ucom.
Proof. intros. apply FullGateSet.decompose_to_cnot_sound. Qed.

Lemma decompose_to_cnot_preserves_WT : forall {dim} (c : circ dim),
  uc_well_typed_l c -> uc_well_typed_l (decompose_to_cnot c).
Proof.
  intros dim c H.
  show_preserves_WT (decompose_to_cnot_preserves_semantics c).
Qed.

Lemma decompose_to_cnot_uses_cnot_gates : forall {dim} (c : circ dim),
  forall_gates only_cnots (decompose_to_cnot c).
Proof. intros. apply FullGateSet.decompose_to_cnot_gates. Qed.

Definition count_I {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App1 U_I _ => true | _ => false end) l).
Definition count_X {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App1 U_X _ => true | _ => false end) l).
Definition count_Y {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App1 U_Y _ => true | _ => false end) l).
Definition count_Z {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App1 U_Z _ => true | _ => false end) l).
Definition count_H {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App1 U_H _ => true | _ => false end) l).
Definition count_S {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App1 U_S _ => true | _ => false end) l).
Definition count_T {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App1 U_T _ => true | _ => false end) l).
Definition count_Sdg {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App1 U_Sdg _ => true | _ => false end) l).
Definition count_Tdg {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App1 U_Tdg _ => true | _ => false end) l).
Definition count_Rx {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App1 (U_Rx _) _ => true | _ => false end) l).
Definition count_Ry {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App1 (U_Ry _) _ => true | _ => false end) l).
Definition count_Rz {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App1 (U_Rz _) _ => true | _ => false end) l).
Definition count_Rzq {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App1 (U_Rzq _) _ => true | _ => false end) l).
Definition count_U1 {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App1 (U_U1 _) _ => true | _ => false end) l).
Definition count_U2 {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App1 (U_U2 _ _) _ => true | _ => false end) l).
Definition count_U3 {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App1 (U_U3 _ _ _) _ => true | _ => false end) l).
Definition count_CX {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App2 U_CX _ _ => true | _ => false end) l).
Definition count_CZ {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App2 U_CZ _ _ => true | _ => false end) l).
Definition count_SWAP {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App2 U_SWAP _ _ => true | _ => false end) l).
Definition count_CCX {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App3 U_CCX _ _ _ => true | _ => false end) l).
Definition count_CCZ {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App3 U_CCZ _ _ _ => true | _ => false end) l).

Definition count_1q {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App1 _ _ => true | _ => false end) l).
Definition count_2q {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App2 _ _ _ => true | _ => false end) l).
Definition count_3q {dim} (l : circ dim) :=
  length (filter (fun g => match g with | App3 _ _ _ _ => true | _ => false end) l).
Definition count_total {dim} (l : circ dim) := length l.

Definition count_rzq_clifford {dim} (l : circ dim) :=
  let f g := match g with
             | App1 (U_Rzq q) _ =>
                 let q' := RzQGateSet.bound q in
                 Qeq_bool q' zero_Q || Qeq_bool q' half_Q ||
                   Qeq_bool q' three_halves_Q || Qeq_bool q' one_Q
             | _ => false end in
  length (filter f l).

Ltac rewrite_count :=
  symmetry; rewrite cons_to_app; rewrite filter_app, app_length; reflexivity.

Lemma count_1q_correct : forall {dim} (l : circ dim),
  count_1q l
    = (count_I l + count_X l + count_Y l + count_Z l +
       count_H l + count_S l + count_T l + count_Sdg l + count_Tdg l +
       count_Rx l + count_Ry l + count_Rz l + count_Rzq l +
       count_U1 l + count_U2 l + count_U3 l)%nat.
Proof.
  intros dim l.
  induction l; simpl.
  reflexivity.
  replace (count_1q (a :: l)) with (count_1q [a] + count_1q l)
    by (unfold count_1q; rewrite_count).
  replace (count_I (a :: l)) with (count_I [a] + count_I l)
    by (unfold count_I; rewrite_count).
  replace (count_X (a :: l)) with (count_X [a] + count_X l)
    by (unfold count_X; rewrite_count).
  replace (count_Y (a :: l)) with (count_Y [a] + count_Y l)
    by (unfold count_Y; rewrite_count).
  replace (count_Z (a :: l)) with (count_Z [a] + count_Z l)
    by (unfold count_Z; rewrite_count).
  replace (count_H (a :: l)) with (count_H [a] + count_H l)
    by (unfold count_H; rewrite_count).
  replace (count_S (a :: l)) with (count_S [a] + count_S l)
    by (unfold count_S; rewrite_count).
  replace (count_T (a :: l)) with (count_T [a] + count_T l)
    by (unfold count_T; rewrite_count).
  replace (count_Sdg (a :: l)) with (count_Sdg [a] + count_Sdg l)
    by (unfold count_Sdg; rewrite_count).
  replace (count_Tdg (a :: l)) with (count_Tdg [a] + count_Tdg l)
    by (unfold count_Tdg; rewrite_count).
  replace (count_Rx (a :: l)) with (count_Rx [a] + count_Rx l)
    by (unfold count_Rx; rewrite_count).
  replace (count_Ry (a :: l)) with (count_Ry [a] + count_Ry l)
    by (unfold count_Ry; rewrite_count).
  replace (count_Rz (a :: l)) with (count_Rz [a] + count_Rz l)
    by (unfold count_Rz; rewrite_count).
  replace (count_Rzq (a :: l)) with (count_Rzq [a] + count_Rzq l)
    by (unfold count_Rzq; rewrite_count).
  replace (count_U1 (a :: l)) with (count_U1 [a] + count_U1 l)
    by (unfold count_U1; rewrite_count).
  replace (count_U2 (a :: l)) with (count_U2 [a] + count_U2 l)
    by (unfold count_U2; rewrite_count).
  replace (count_U3 (a :: l)) with (count_U3 [a] + count_U3 l)
    by (unfold count_U3; rewrite_count).
  rewrite IHl. clear.
  repeat rewrite Nat.add_assoc.
  repeat rewrite (Nat.add_comm _ (_ [a])).
  repeat rewrite Nat.add_assoc.
  do 16 (apply f_equal2; auto).
  destruct a; dependent destruction f; reflexivity.
Qed.

Lemma count_2q_correct : forall {dim} (l : circ dim),
  count_2q l
    = (count_CX l + count_CZ l + count_SWAP l)%nat.
Proof.
  intros dim l.
  induction l; simpl.
  reflexivity.
  replace (count_2q (a :: l)) with (count_2q [a] + count_2q l)
    by (unfold count_2q; rewrite_count).
  replace (count_CX (a :: l)) with (count_CX [a] + count_CX l)
    by (unfold count_CX; rewrite_count).
  replace (count_CZ (a :: l)) with (count_CZ [a] + count_CZ l)
    by (unfold count_CZ; rewrite_count).
  replace (count_SWAP (a :: l)) with (count_SWAP [a] + count_SWAP l)
    by (unfold count_SWAP; rewrite_count).
  rewrite IHl. clear.
  repeat rewrite Nat.add_assoc.
  repeat rewrite (Nat.add_comm _ (_ [a])).
  repeat rewrite Nat.add_assoc.
  do 3 (apply f_equal2; auto).
  destruct a; dependent destruction f; reflexivity.
Qed.

Lemma count_3q_correct : forall {dim} (l : circ dim),
  count_3q l
    = (count_CCX l + count_CCZ l)%nat.
Proof.
  intros dim l.
  induction l; simpl.
  reflexivity.
  replace (count_3q (a :: l)) with (count_3q [a] + count_3q l)
    by (unfold count_3q; rewrite_count).
  replace (count_CCX (a :: l)) with (count_CCX [a] + count_CCX l)
    by (unfold count_CCX; rewrite_count).
  replace (count_CCZ (a :: l)) with (count_CCZ [a] + count_CCZ l)
    by (unfold count_CCZ; rewrite_count).
  rewrite IHl. clear.
  repeat rewrite Nat.add_assoc.
  repeat rewrite (Nat.add_comm _ (_ [a])).
  repeat rewrite Nat.add_assoc.
  do 2 (apply f_equal2; auto).
  destruct a; dependent destruction f; reflexivity.
Qed.

Lemma count_total_correct : forall {dim} (l : circ dim),
  count_total l = (count_1q l + count_2q l + count_3q l)%nat.
Proof.
  intros dim l.
  induction l; simpl.
  reflexivity.
  replace (count_1q (a :: l)) with (count_1q [a] + count_1q l)
    by (unfold count_1q; rewrite_count).
  replace (count_2q (a :: l)) with (count_2q [a] + count_2q l)
    by (unfold count_2q; rewrite_count).
  replace (count_3q (a :: l)) with (count_3q [a] + count_3q l)
    by (unfold count_3q; rewrite_count).
  rewrite IHl. clear.
  destruct a; dependent destruction f; simpl; lia.
Qed.

(** * IBM gate set optimizations **)

Definition optimize_ibm {dim} (c : circ dim) : circ dim :=
  IBM_to_full (Optimize1qGates.optimize_1q_gates (full_to_IBM c)).

Lemma optimize_ibm_preserves_semantics : forall {dim} (c : circ dim),
  uc_well_typed_l c -> (optimize_ibm c ≅l≅ c)%ucom.
Proof.
  intros dim c H.
  unfold optimize_ibm.
  erewrite IBM_to_full_cong.
  apply uc_equiv_cong_l.
  apply IBM_to_full_inv.
  apply Optimize1qGates.optimize_1q_gates_sound.
  apply FullGateSet.full_to_IBM_WT.
  assumption.
Qed.

Lemma optimize_ibm_preserves_WT : forall {dim} (c : circ dim),
  uc_well_typed_l c -> uc_well_typed_l (optimize_ibm c).
Proof.
  intros dim c H.
  show_preserves_WT_cong (optimize_ibm_preserves_semantics c H).
Qed.

Lemma optimize_ibm_preserves_mapping : forall {dim} (c : circ dim) (cg : c_graph),
  respects_constraints_directed (is_in_graph cg) U_CX c ->
  respects_constraints_directed (is_in_graph cg) U_CX (optimize_ibm c).
Proof.
  intros.
  apply IBM_to_full_preserves_mapping.
  apply Optimize1qGates.optimize_1q_gates_respects_constraints.
  apply full_to_IBM_preserves_mapping.
  assumption.
Qed.

(** * RzQ gate set optimizations **)

Definition not_propagation {dim} (c : circ dim) : circ dim :=
  RzQ_to_full (NotPropagation.not_propagation (full_to_RzQ c)).

Definition hadamard_reduction {dim} (c : circ dim) : circ dim :=
  RzQ_to_full (HadamardReduction.hadamard_reduction (full_to_RzQ c)).

Definition cancel_single_qubit_gates {dim} (c : circ dim) : circ dim :=
  RzQ_to_full (GateCancellation.cancel_single_qubit_gates (full_to_RzQ c)).

Definition cancel_two_qubit_gates {dim} (c : circ dim) : circ dim :=
  RzQ_to_full (GateCancellation.cancel_two_qubit_gates (full_to_RzQ c)).

Definition merge_rotations {dim} (c : circ dim) : circ dim :=
  RzQ_to_full (RotationMerging.merge_rotations (full_to_RzQ c)).

(* optimize_nam function applies our optimizations in the following order,
   as designed by Nam et al. :
   0, 1, 3, 2, 3, 1, 2, 4, 3, 2

   0 - not propagation
   1 - hadamard reduction
   2 - single qubit gate cancellation
   3 - two qubit gate cancellation
   4 - rotation merging *)

Definition optimize_nam {dim} (c : circ dim) : circ dim :=
  RzQ_to_full
    (GateCancellation.cancel_single_qubit_gates
      (GateCancellation.cancel_two_qubit_gates
        (RotationMerging.merge_rotations
          (GateCancellation.cancel_single_qubit_gates
            (HadamardReduction.hadamard_reduction
              (GateCancellation.cancel_two_qubit_gates
                (GateCancellation.cancel_single_qubit_gates
                  (GateCancellation.cancel_two_qubit_gates
                    (HadamardReduction.hadamard_reduction
                      (NotPropagation.not_propagation
                        (full_to_RzQ c))))))))))).

(* Light version of the optimizer that excludes rotation merging
   (used for evaluating on QFT & adder programs). *)
Definition optimize_nam_light {dim} (c : circ dim) : circ dim :=
  RzQ_to_full
    (GateCancellation.cancel_single_qubit_gates
      (HadamardReduction.hadamard_reduction
        (GateCancellation.cancel_two_qubit_gates
          (GateCancellation.cancel_single_qubit_gates
            (GateCancellation.cancel_two_qubit_gates
              (HadamardReduction.hadamard_reduction
                (NotPropagation.not_propagation
                  (full_to_RzQ c)))))))).

(* LCR optimizer for multiple iterations. *)
Definition optimize_nam_lcr {dim} (c : circ dim) : option (circ dim * circ dim * circ dim) :=
  LCR c optimize_nam (fun n => @match_gate n).

Lemma cancel_single_qubit_gates_sound' : forall {dim} (l : RzQ_ucom_l dim),
  uc_well_typed_l l -> RzQList.uc_cong_l (GateCancellation.cancel_single_qubit_gates l) l.
Proof.
  intros. apply RzQList.uc_equiv_cong_l.
  apply GateCancellation.cancel_single_qubit_gates_sound. assumption.
Qed.

Lemma cancel_two_qubit_gates_sound' : forall {dim} (l : RzQ_ucom_l dim),
  uc_well_typed_l l -> RzQList.uc_cong_l (GateCancellation.cancel_two_qubit_gates l) l.
Proof.
  intros. apply RzQList.uc_equiv_cong_l.
  apply GateCancellation.cancel_two_qubit_gates_sound. assumption.
Qed.

Lemma merge_rotations_sound' : forall {dim} (l : RzQ_ucom_l dim),
  uc_well_typed_l l -> RzQList.uc_cong_l (RotationMerging.merge_rotations l) l.
Proof.
  intros. apply RzQList.uc_equiv_cong_l.
  apply RotationMerging.merge_rotations_sound. assumption.
Qed.

Ltac show_preserves_semantics_nam :=
  unfold not_propagation, hadamard_reduction, cancel_single_qubit_gates, cancel_two_qubit_gates, merge_rotations, optimize_nam, optimize_nam_light;
  erewrite RzQ_to_full_cong;
  [ apply RzQ_to_full_inv
  | repeat (try rewrite NotPropagation.not_propagation_sound;
            try rewrite HadamardReduction.hadamard_reduction_sound;
            try rewrite cancel_single_qubit_gates_sound';
            try rewrite cancel_two_qubit_gates_sound';
            try rewrite merge_rotations_sound';
            try apply FullGateSet.full_to_RzQ_WT;
            try apply NotPropagation.not_propagation_WT;
            try apply HadamardReduction.hadamard_reduction_WT;
            try apply GateCancellation.cancel_single_qubit_gates_WT;
            try apply GateCancellation.cancel_two_qubit_gates_WT;
            try apply RotationMerging.merge_rotations_WT;
            try assumption; try reflexivity) ].

Lemma not_propagation_preserves_semantics : forall {dim} (c : circ dim),
  uc_well_typed_l c -> (not_propagation c ≅l≅ c)%ucom.
Proof. intros. show_preserves_semantics_nam. Qed.

Lemma not_propagation_preserves_WT : forall {dim} (c : circ dim),
  uc_well_typed_l c -> uc_well_typed_l (not_propagation c).
Proof.
  intros dim c H.
  show_preserves_WT_cong (not_propagation_preserves_semantics c H).
Qed.

Ltac show_preserves_mapping_nam :=
  unfold not_propagation, hadamard_reduction, cancel_single_qubit_gates, cancel_two_qubit_gates, merge_rotations, optimize_nam, optimize_nam_light;
  repeat (try apply RzQ_to_full_preserves_mapping;
          try apply NotPropagation.not_propagation_respects_constraints;
          try apply HadamardReduction.hadamard_reduction_respects_constraints;
          try apply GateCancellation.cancel_single_qubit_gates_respects_constraints;
          try apply GateCancellation.cancel_two_qubit_gates_respects_constraints;
          try apply RotationMerging.merge_rotations_respects_constraints;
          try apply full_to_RzQ_preserves_mapping;
          try assumption).

Lemma not_propagation_preserves_mapping : forall {dim} (c : circ dim) (cg : c_graph),
  respects_constraints_directed (is_in_graph cg) U_CX c ->
  respects_constraints_directed (is_in_graph cg) U_CX (not_propagation c).
Proof. intros. show_preserves_mapping_nam. Qed.

Lemma hadamard_reduction_preserves_semantics : forall {dim} (c : circ dim),
  (hadamard_reduction c ≅l≅ c)%ucom.
Proof. intros. show_preserves_semantics_nam. Qed.

Lemma hadamard_reduction_preserves_WT : forall {dim} (c : circ dim),
  uc_well_typed_l c -> uc_well_typed_l (hadamard_reduction c).
Proof.
  intros dim c H.
  show_preserves_WT_cong (hadamard_reduction_preserves_semantics c).
Qed.

Lemma hadamard_reduction_preserves_mapping : forall {dim} (c : circ dim) (cg : c_graph),
  respects_constraints_directed (is_in_graph cg) U_CX c ->
  respects_constraints_directed (is_in_graph cg) U_CX (hadamard_reduction c).
Proof. intros. show_preserves_mapping_nam. Qed.

Lemma cancel_single_qubit_gates_preserves_semantics : forall {dim} (c : circ dim),
  uc_well_typed_l c -> (cancel_single_qubit_gates c ≅l≅ c)%ucom.
Proof. intros. show_preserves_semantics_nam. Qed.

Lemma cancel_single_qubit_gates_preserves_WT : forall {dim} (c : circ dim),
  uc_well_typed_l c -> uc_well_typed_l (cancel_single_qubit_gates c).
Proof.
  intros dim c H.
  show_preserves_WT_cong (cancel_single_qubit_gates_preserves_semantics c H).
Qed.

Lemma cancel_single_qubit_gates_preserves_mapping : forall {dim} (c : circ dim) (cg : c_graph),
  respects_constraints_directed (is_in_graph cg) U_CX c ->
  respects_constraints_directed (is_in_graph cg) U_CX (cancel_single_qubit_gates c).
Proof. intros. show_preserves_mapping_nam. Qed.

Lemma cancel_two_qubit_gates_preserves_semantics : forall {dim} (c : circ dim),
  uc_well_typed_l c -> (cancel_two_qubit_gates c ≅l≅ c)%ucom.
Proof. intros. show_preserves_semantics_nam. Qed.

Lemma cancel_two_qubit_gates_preserves_WT : forall {dim} (c : circ dim),
  uc_well_typed_l c -> uc_well_typed_l (cancel_two_qubit_gates c).
Proof.
  intros dim c H.
  show_preserves_WT_cong (cancel_two_qubit_gates_preserves_semantics c H).
Qed.

Lemma cancel_two_qubit_gates_preserves_mapping : forall {dim} (c : circ dim) (cg : c_graph),
  respects_constraints_directed (is_in_graph cg) U_CX c ->
  respects_constraints_directed (is_in_graph cg) U_CX (cancel_two_qubit_gates c).
Proof. intros. show_preserves_mapping_nam. Qed.

Lemma merge_rotations_preserves_semantics : forall {dim} (c : circ dim),
  uc_well_typed_l c -> (merge_rotations c ≅l≅ c)%ucom.
Proof. intros. show_preserves_semantics_nam. Qed.

Lemma merge_rotations_preserves_WT : forall {dim} (c : circ dim),
  uc_well_typed_l c -> uc_well_typed_l (merge_rotations c).
Proof.
  intros dim c H.
  show_preserves_WT_cong (merge_rotations_preserves_semantics c H).
Qed.

Lemma merge_rotations_preserves_mapping : forall {dim} (c : circ dim) (cg : c_graph),
  respects_constraints_directed (is_in_graph cg) U_CX c ->
  respects_constraints_directed (is_in_graph cg) U_CX (merge_rotations c).
Proof. intros. show_preserves_mapping_nam. Qed.

Lemma optimize_nam_preserves_semantics : forall {dim} (c : circ dim),
  uc_well_typed_l c -> (optimize_nam c ≅l≅ c)%ucom.
Proof. intros. show_preserves_semantics_nam. Qed.

Lemma optimize_nam_preserves_WT : forall {dim} (c : circ dim),
  uc_well_typed_l c -> uc_well_typed_l (optimize_nam c).
Proof.
  intros dim c H.
  show_preserves_WT_cong (optimize_nam_preserves_semantics c H).
Qed.

Lemma optimize_nam_preserves_mapping : forall {dim} (c : circ dim) (cg : c_graph),
  respects_constraints_directed (is_in_graph cg) U_CX c ->
  respects_constraints_directed (is_in_graph cg) U_CX (optimize_nam c).
Proof. intros. show_preserves_mapping_nam. Qed.

Lemma optimize_nam_light_preserves_semantics : forall {dim} (c : circ dim),
  uc_well_typed_l c -> (optimize_nam_light c ≅l≅ c)%ucom.
Proof. intros. show_preserves_semantics_nam. Qed.

Lemma optimize_nam_light_preserves_WT : forall {dim} (c : circ dim),
  uc_well_typed_l c -> uc_well_typed_l (optimize_nam_light c).
Proof.
  intros dim c H.
  show_preserves_WT_cong (optimize_nam_light_preserves_semantics c H).
Qed.

Lemma optimize_nam_light_preserves_mapping : forall {dim} (c : circ dim) (cg : c_graph),
  respects_constraints_directed (is_in_graph cg) U_CX c ->
  respects_constraints_directed (is_in_graph cg) U_CX (optimize_nam_light c).
Proof. intros. show_preserves_mapping_nam. Qed.

Lemma optimize_nam_lcr_preserves_semantics : forall {dim} (c0 l c r : circ dim) n,
  n > 2 -> uc_well_typed_l c0 ->
  optimize_nam_lcr c0 = Some (l, c, r) ->
  (niter c0 n ≅l≅ (l ++ (niter c (n - 2)) ++ r))%ucom.
Proof.
  intros dim c0 l c r n Hn WT H.
  eapply LCR_correct in H.
  apply H.
  all: try assumption.
  apply optimize_nam_preserves_semantics.
  apply optimize_nam_preserves_WT.
Qed.

Lemma niter_WT : forall {dim} (c : circ dim) n,
  uc_well_typed_l c -> uc_well_typed_l (niter c n).
Proof.
  intros dim c n WT.
  induction n.
  constructor.
  eapply uc_well_typed_l_implies_dim_nonzero.
  apply WT.
  simpl.
  apply uc_well_typed_l_app; split; assumption.
Qed.

Lemma niter_WT_inv : forall {dim} (c : circ dim) n,
  n > 0 -> uc_well_typed_l (niter c n) -> uc_well_typed_l c.
Proof.
  intros dim c n Hn WT.
  destruct n; try lia.
  induction n; simpl in WT.
  rewrite app_nil_r in WT.
  assumption.
  apply IHn; try lia.
  simpl.
  apply uc_well_typed_l_app in WT as [_ WT].
  assumption.
Qed.

Lemma optimize_nam_lcr_preserves_WT : forall {dim} (c0 l c r : circ dim) n,
  n > 2 -> uc_well_typed_l c0 ->
  optimize_nam_lcr c0 = Some (l, c, r) ->
  uc_well_typed_l l /\ uc_well_typed_l c /\ uc_well_typed_l r.
Proof.
  intros dim c0 l c r n Hn WT H.
  eapply optimize_nam_lcr_preserves_semantics in H; try apply Hn; auto.
  apply uc_cong_l_implies_WT in H.
  apply uc_well_typed_l_app in H as [H1 H23].
  apply uc_well_typed_l_app in H23 as [H2 H3].
  repeat split; try assumption.
  eapply niter_WT_inv; try apply H2. lia.
  apply niter_WT.
  assumption.
Qed.

Lemma optimize_nam_lcr_preserves_mapping : forall {dim} (c0 l c r : circ dim) (cg : c_graph),
  respects_constraints_directed (is_in_graph cg) U_CX c0 ->
  optimize_nam_lcr c0 = Some (l, c, r) ->
  respects_constraints_directed (is_in_graph cg) U_CX l
    /\ respects_constraints_directed (is_in_graph cg) U_CX c
    /\ respects_constraints_directed (is_in_graph cg) U_CX r.
Proof.
  intros dim c0 l c r cg Hcg H.
  eapply MappingConstraints.LCR_respects_constraints in H as [H0 [H1 H2]].
  repeat split.
  apply H0. apply H2. apply H1.
  intros.
  apply optimize_nam_preserves_mapping.
  assumption.
  assumption.
Qed.

(** * Full 'optimize' function *)

Definition optimize {dim} (c : circ dim) : circ dim :=
  optimize_ibm (optimize_nam c).

Lemma optimize_preserves_semantics : forall {dim} (c : circ dim),
  uc_well_typed_l c -> (optimize c ≅l≅ c)%ucom.
Proof.
  intros dim c H.
  unfold optimize.
  rewrite optimize_ibm_preserves_semantics.
  apply optimize_nam_preserves_semantics.
  assumption.
  apply optimize_nam_preserves_WT.
  assumption.
Qed.

Lemma optimize_preserves_WT : forall {dim} (c : circ dim),
  uc_well_typed_l c -> uc_well_typed_l (optimize c).
Proof.
  intros dim c H.
  show_preserves_WT_cong (optimize_preserves_semantics c H).
Qed.

Lemma optimize_preserves_mapping : forall {dim} (c : circ dim) (cg : c_graph),
  respects_constraints_directed (is_in_graph cg) U_CX c ->
  respects_constraints_directed (is_in_graph cg) U_CX (optimize c).
Proof.
  intros.
  apply optimize_ibm_preserves_mapping.
  apply optimize_nam_preserves_mapping.
  assumption.
Qed.

(** * Circuit mapping **)

Definition swap_route {dim} (c : circ dim) (lay : layout) (cg : c_graph) (get_path : path_finding_fun) :=
  let n := graph_dim cg in
  let (c,_) := SwapRoute.swap_route (full_to_map (cast c n)) lay get_path in
  map_to_full c.

Definition decompose_swaps {dim} (c : circ dim) (cg : c_graph) :=
  map_to_full (SwapRoute.decompose_swaps_and_cnots (full_to_map c) (is_in_graph cg)).

Definition trivial_layout n : layout := Layouts.trivial_layout n.
Definition check_list l : bool := Layouts.check_list l.
Definition list_to_layout l : layout := Layouts.list_to_layout l.
Definition layout_to_list (lay : layout) n : list nat :=
  map (fun ox => match ox with Some x => x | _ => O end) (Layouts.layout_to_list n lay).
Definition greedy_layout {dim} (c : circ dim) (cg : c_graph) (q_ordering : option nat -> list nat) : layout :=
  let n := graph_dim cg in
  GreedyLayout.greedy_layout (full_to_map (cast c n)) n q_ordering.

Definition beq_tup t t' :=
  match t, t' with
  | (n1, n2), (n1', n2') => (n1 =? n1') && (n2 =? n2')
  end.

Definition make_lnn n : c_graph := (n, LNN.is_in_graph n).
Definition make_lnn_ring n : c_graph := (n, LNNRing.is_in_graph n).
Definition make_grid m n : c_graph := (m * n, Grid.is_in_graph m n).

Definition c_graph_from_coupling_map (n : nat) (cmap : list (nat * nat)) : c_graph :=
  (n, fun n1 n2 => existsb (beq_tup (n1, n2)) cmap).

Definition lnn_path_finding_fun (n : nat) : path_finding_fun := LNN.get_path.
Definition lnn_ring_path_finding_fun n : path_finding_fun := LNNRing.get_path n.
Definition grid_path_finding_fun (m n : nat) : path_finding_fun := Grid.get_path n.

Definition lnn_qubit_ordering_fun n : qubit_ordering_fun := LNN.q_ordering n.
Definition lnn_ring_qubit_ordering_fun n : qubit_ordering_fun := LNNRing.q_ordering n.

Definition get_path_valid (cg : c_graph) (get_path : path_finding_fun) :=
  ConnectivityGraph.get_path_valid (fst cg) get_path (snd cg).

Lemma lnn_path_finding_fun_valid : forall n,
  get_path_valid (make_lnn n) (lnn_path_finding_fun n).
Proof. intros. apply LNN.lnn_get_path_valid. Qed.

Lemma lnn_ring_path_finding_fun_valid : forall n,
  get_path_valid (make_lnn_ring n) (lnn_ring_path_finding_fun n).
Proof. intros. apply LNNRing.lnn_ring_get_path_valid. Qed.

Lemma grid_path_finding_fun_valid : forall m n,
  get_path_valid (make_grid m n) (grid_path_finding_fun m n).
Proof.
  intros.
  intros ? ? ? ? ?.
  apply Grid.get_path_valid; auto.
Qed.

Lemma lnn_qubit_ordering_fun_valid : forall n,
  valid_q_ordering (lnn_qubit_ordering_fun n) n.
Proof. intros. apply LNN.lnn_q_ordering_valid. Qed.

Lemma lnn_ring_qubit_ordering_fun_valid : forall n,
  valid_q_ordering (lnn_ring_qubit_ordering_fun n) n.
Proof. intros. apply LNNRing.lnn_ring_q_ordering_valid. Qed.

Module MVP := MappingValidationProofs FullGateSet.

Lemma list_to_ucom_map_to_full : forall {dim} (l : gate_list _ dim),
  uc_equiv (MVP.SRP.MapList.list_to_ucom l) (FullList.list_to_ucom (map_to_full l)).
Proof.
  intros.
  induction l.
  reflexivity.
  simpl.
  unfold map_to_full.
  rewrite change_gate_set_cons.
  rewrite FullList.list_to_ucom_append.
  destruct a; rewrite IHl; apply useq_mor; try reflexivity.
  all: dependent destruction u; simpl; rewrite SKIP_id_r; reflexivity.
Qed.

Lemma list_to_ucom_full_to_map : forall {dim} (l : circ dim),
  uc_equiv (FullList.list_to_ucom l) (MVP.SRP.MapList.list_to_ucom (full_to_map l)).
Proof.
  intros.
  induction l.
  reflexivity.
  simpl.
  unfold full_to_map.
  rewrite change_gate_set_cons.
  rewrite MapList.list_to_ucom_append.
  destruct a; rewrite IHl; apply useq_mor; try reflexivity.
  all: dependent destruction f; simpl; rewrite SKIP_id_r; try reflexivity.
  all: repeat rewrite <- useq_assoc; reflexivity.
Qed.

Lemma swap_route_preserves_semantics : forall {dim} (c : circ dim) (lay : layout) (cg : c_graph) (get_path : path_finding_fun),
  let n := graph_dim cg in
  uc_well_typed_l (cast c n) ->
  layout_bijective n lay ->
  get_path_valid cg get_path ->
  cast c n ≡x swap_route c lay cg get_path.
Proof.
  intros dim c lay cg get_path n WT WF Hpath.
  subst n.
  unfold swap_route.
  destruct (SwapRoute.swap_route (full_to_map (cast c (graph_dim cg))) lay get_path) eqn:sr.
  assert (srWF:=sr).
  apply MVP.SRP.swap_route_WF in srWF; auto.
  apply MVP.SRP.swap_route_sound in sr; auto.
  unfold MVP.SRP.uc_equiv_perm_ex in sr.
  unfold uc_equiv_perm.
  exists (get_log lay). exists (get_phys l).
  repeat split.
  apply get_log_perm. assumption.
  apply get_phys_perm. assumption.
  unfold eval, MVP.SRP.MapList.eval in *.
  rewrite <- list_to_ucom_map_to_full, <- sr.
  rewrite list_to_ucom_full_to_map.
  reflexivity.
  apply full_to_map_WT. assumption.
  intros n1 n2 Hn1 Hn2 Hneq.
  destruct (Hpath n1 n2 Hn1 Hn2 Hneq) as [_ [_ [_ [H _]]]].
  apply H.
  apply full_to_map_WT. assumption.
  intros n1 n2 Hn1 Hn2 Hneq.
  destruct (Hpath n1 n2 Hn1 Hn2 Hneq) as [_ [_ [_ [H _]]]].
  apply H.
Qed.

Lemma swap_route_preserves_WT : forall {dim} (c : circ dim) (lay : layout) (cg : c_graph) (get_path : path_finding_fun),
  let n := graph_dim cg in
  uc_well_typed_l (cast c n) ->
  layout_bijective n lay ->
  get_path_valid cg get_path ->
  uc_well_typed_l (swap_route c lay cg get_path).
Proof.
  intros dim c lay cg get_path n WT WF Hpath.
  specialize (swap_route_preserves_semantics _ _ _ get_path WT WF Hpath) as H.
  destruct H as [p1 [p2 [Hp1 [Hp2 H]]]].
  apply list_to_ucom_WT.
  apply uc_eval_nonzero_iff.
  apply list_to_ucom_WT in WT.
  apply uc_eval_nonzero_iff in WT.
  intro contra.
  unfold eval in H.
  rewrite contra in H.
  rewrite Mmult_0_r, Mmult_0_l in H.
  contradiction.
Qed.

Lemma swap_route_respects_constraints_undirected : forall {dim} (c : circ dim) (lay : layout) (cg : c_graph) (get_path : path_finding_fun),
  let n := graph_dim cg in
  uc_well_typed_l (cast c n) ->
  layout_bijective n lay ->
  get_path_valid cg get_path ->
  respects_constraints_undirected (is_in_graph cg) (swap_route c lay cg get_path).
Proof.
  intros dim c lay cg get_path n WT WF Hpath.
  subst n.
  unfold swap_route.
  destruct (SwapRoute.swap_route (full_to_map (cast c (graph_dim cg))) lay get_path) eqn:sr.
  apply MVP.SRP.swap_route_respects_undirected with (is_in_graph:=is_in_graph cg) in sr; auto.
  apply map_to_full_preserves_mapping_undirected. assumption.
  apply full_to_map_WT. assumption.
Qed.

Lemma map_to_full_equiv : forall {dim} (l l' : gate_list _ dim),
  MVP.SRP.MapList.uc_equiv_l l l' ->
  uc_equiv_l (map_to_full l) (map_to_full l').
Proof.
  intros dim l l' H.
  unfold uc_equiv_l.
  unfold MVP.SRP.MapList.uc_equiv_l in H.
  rewrite <- 2 list_to_ucom_map_to_full.
  assumption.
Qed.

Lemma decompose_swaps_preserves_semantics : forall {dim} (c : circ dim) (cg : c_graph),
  uc_equiv_l (decompose_swaps c cg) c.
Proof.
  intros.
  unfold decompose_swaps.
  erewrite map_to_full_equiv.
  apply map_to_full_inv.
  apply decompose_swaps_and_cnots_sound.
Qed.

Lemma decompose_swaps_preserves_WT : forall {dim} (c : circ dim) (cg : c_graph),
  uc_well_typed_l c ->
  uc_well_typed_l (decompose_swaps c cg).
Proof.
  intros dim c cg WT.
  specialize (decompose_swaps_preserves_semantics c cg) as H.
  apply list_to_ucom_WT.
  apply uc_eval_nonzero_iff.
  apply list_to_ucom_WT in WT.
  apply uc_eval_nonzero_iff in WT.
  intro contra.
  unfold uc_equiv_l, uc_equiv in H.
  rewrite contra in H.
  rewrite H in WT.
  contradiction.
Qed.

Lemma decompose_swaps_respects_constraints : forall {dim} (c : circ dim) (cg : c_graph),
  respects_constraints_undirected (is_in_graph cg) c ->
  respects_constraints_directed (is_in_graph cg) U_CX (decompose_swaps c cg).
Proof.
  intros.
  unfold decompose_swaps.
  apply map_to_full_preserves_mapping_directed.
  apply decompose_swaps_and_cnots_respects_directed.
  apply full_to_map_preserves_mapping_undirected.
  assumption.
Qed.

Lemma trivial_layout_well_formed : forall n, layout_bijective n (trivial_layout n).
Proof. intros. apply Layouts.trivial_layout_bijective. Qed.

Lemma list_to_layout_well_formed : forall l,
  check_list l = true -> layout_bijective (length l) (list_to_layout l).
Proof. intros l H. apply Layouts.check_list_layout_bijective. auto. Qed.

Lemma greedy_layout_well_formed : forall {dim} (c : circ dim) (cg : c_graph) (q_ordering : qubit_ordering_fun),
  let n := graph_dim cg in
  uc_well_typed_l (cast c n) ->
  valid_q_ordering q_ordering (graph_dim cg) ->
  layout_bijective n (greedy_layout c cg q_ordering).
Proof.
  intros.
  apply GreedyLayout.greedy_layout_bijective.
  apply full_to_map_WT.
  assumption.
  assumption.
Qed.

(** * Mapping validation **)

Definition remove_swaps {dim} (c : circ dim) (lay : layout) :=
  let (c,_) := MappingValidation.remove_swaps (full_to_map c) lay in
  map_to_full c.

Definition check_swap_equivalence {dim} (c1 c2 : circ dim) (lay1 lay2 : layout) :=
  MappingValidation.is_swap_equivalent (full_to_map c1) (full_to_map c2) lay1 lay2
    (fun n => @MappingGateSet.match_gate (FullGateSet.U 1) n FullGateSet.match_gate).

Definition check_constraints {dim} (c : circ dim) (cg : c_graph) :=
  MappingValidation.check_constraints (full_to_map c) (is_in_graph cg).

Lemma full_to_map_inv : forall {dim} (l : _ dim),
  MVP.SRP.MapList.uc_equiv_l (full_to_map (map_to_full l)) l.
Proof.
  intros dim l.
  induction l.
  reflexivity.
  unfold full_to_map, map_to_full.
  rewrite change_gate_set_cons.
  rewrite change_gate_set_app.
  rewrite IHl.
  rewrite cons_to_app.
  MVP.SRP.MapList.apply_app_congruence.
  destruct a; dependent destruction m;
  unfold change_gate_set; simpl; reflexivity.
Qed.

Lemma remove_swaps_preserves_semantics : forall {dim} (c : circ dim) (lay : layout),
  uc_well_typed_l c ->
  layout_bijective dim lay ->
  remove_swaps c lay ≡x c.
Proof.
  intros dim c lay WT WF.
  unfold remove_swaps.
  destruct (MappingValidation.remove_swaps (full_to_map c) lay) eqn:rs.
  assert (rsWF:=rs).
  apply MVP.remove_swaps_WF in rsWF; auto.
  apply MVP.remove_swaps_sound in rs; auto.
  unfold MVP.SRP.uc_equiv_perm_ex in rs.
  symmetry.
  unfold uc_equiv_perm.
  exists (get_phys lay). exists (get_log l).
  repeat split.
  apply get_phys_perm. assumption.
  apply get_log_perm. assumption.
  unfold eval, MVP.SRP.MapList.eval in *.
  rewrite <- list_to_ucom_full_to_map in rs.
  rewrite rs.
  apply f_equal2; try reflexivity.
  apply f_equal2; try reflexivity.
  rewrite list_to_ucom_full_to_map.
  rewrite full_to_map_inv.
  reflexivity.
  apply full_to_map_WT. assumption.
  apply full_to_map_WT. assumption.
Qed.

Lemma remove_swaps_preserves_WT : forall {dim} (c : circ dim) (lay : layout),
  uc_well_typed_l c ->
  layout_bijective dim lay ->
  uc_well_typed_l (remove_swaps c lay).
Proof.
  intros dim c lay WT WF.
  specialize (remove_swaps_preserves_semantics c lay WT WF) as H.
  symmetry in H.
  destruct H as [p1 [p2 [Hp1 [Hp2 H]]]].
  apply list_to_ucom_WT.
  apply uc_eval_nonzero_iff.
  apply list_to_ucom_WT in WT.
  apply uc_eval_nonzero_iff in WT.
  intro contra.
  unfold eval in H.
  rewrite contra in H.
  rewrite Mmult_0_r, Mmult_0_l in H.
  contradiction.
Qed.

Lemma check_swap_equivalence_correct : forall dim (c1 c2 : circ dim) (lay1 lay2 : layout),
  uc_well_typed_l c1 ->
  uc_well_typed_l c2 ->
  layout_bijective dim lay1 ->
  layout_bijective dim lay2 ->
  check_swap_equivalence c1 c2 lay1 lay2 = true ->
  c1 ≡x c2.
Proof.
  intros dim c1 c2 lay1 lay2 WT1 WT2 WF1 WF2 H.
  unfold check_swap_equivalence in H.
  unfold is_swap_equivalent in H.
  destruct (MappingValidation.check_swap_equivalence (full_to_map c1)
                                                     (full_to_map c2) lay1 lay2
                                                     (fun n : nat => MappingGateSet.match_gate match_gate)) eqn:mv.
  assert (mvWF:=mv).
  destruct p.
  2: inversion H.
  apply MVP.check_swap_equivalence_implies_equivalence in mv; auto.
  apply MVP.check_swap_equivalence_layouts_WF in mvWF as [? ?]; auto.
  unfold MVP.SRP.uc_equiv_perm_ex in mv.
  exists (get_phys lay1 ∘ get_log lay2)%prg.
  exists (get_phys l0 ∘ get_log l)%prg.
  repeat split.
  apply Permutations.permutation_compose.
  apply get_phys_perm; auto.
  apply get_log_perm; auto.
  apply Permutations.permutation_compose.
  apply get_phys_perm; auto.
  apply get_log_perm; auto.
  unfold eval.
  unfold MVP.SRP.MapList.eval in mv.
  rewrite <- 2 list_to_ucom_full_to_map in mv.
  apply mv.
  all: apply full_to_map_WT; assumption.
Qed.

Lemma check_constraints_correct : forall dim (c : circ dim) (cg : c_graph),
  check_constraints c cg = true ->
  respects_constraints_directed (is_in_graph cg) MappingGateSet.UMap_CNOT (full_to_map c).
Proof. intros. apply MVP.check_constraints_implies_respect_constraints. auto. Qed.

(** * Example verified composition of transformations **)

Definition optimize_and_map_to_lnn_ring_16 {dim} (c : circ dim) :=
  let cg := make_lnn_ring 16 in
  let get_path := lnn_ring_path_finding_fun 16 in
  let q_ordering := lnn_ring_qubit_ordering_fun 16 in
  if check_well_typed c 16
  then
    let c1 := optimize_nam c in                 (* optimization #1 *)
    let lay := greedy_layout c cg q_ordering in
    let c2 := swap_route c1 lay cg get_path in  (* mapping *)
    let c3 := decompose_swaps c2 cg in          (* optimized SWAP decomposition *)
    Some (optimize c3)                          (* optimization #2 *)
  else None.

Lemma cast_same : forall {dim} (c : circ dim), cast c dim = c.
Proof.
  intros dim c.
  induction c.
  reflexivity.
  simpl.
  destruct a; rewrite IHc; reflexivity.
Qed.

Lemma optimize_and_map_to_lnn_ring_16_preserves_semantics : forall (c : circ _) c',
  optimize_and_map_to_lnn_ring_16 c = Some c' ->
  c ≅x c'.
Proof.
  intros c c' H.
  unfold optimize_and_map_to_lnn_ring_16 in H.
  remember (make_lnn_ring 16) as cg.
  remember (lnn_ring_path_finding_fun 16) as get_path.
  remember (greedy_layout c cg (lnn_ring_qubit_ordering_fun 16)) as lay.
  assert (Hpath : get_path_valid cg get_path).
  { subst. apply lnn_ring_path_finding_fun_valid. }
  clear Heqget_path.
  destruct (check_well_typed c 16) eqn:WT; inversion H.
  apply check_well_typed_correct in WT.
  replace 16 with (graph_dim cg) in *.
  rewrite cast_same in WT.
  assert (WF : layout_bijective (graph_dim cg) lay).
  { subst. apply greedy_layout_well_formed.
    rewrite cast_same. assumption.
    apply lnn_ring_qubit_ordering_fun_valid. }
  clear - WT Hpath WF.
  specialize (swap_route_preserves_semantics (optimize_nam c) lay cg get_path) as Hmap.
  destruct Hmap as [p1 [p2 [Hp1 [Hp2 Hmap]]]]; auto.
  rewrite cast_same.
  apply optimize_nam_preserves_WT. auto.
  exists p1. exists p2.
  repeat split; auto.
  rewrite cast_same in Hmap.
  specialize (optimize_nam_preserves_semantics c WT) as Hnam.
  destruct Hnam as [x Hnam].
  assert (Haux: eval c = Cexp (- x) .* eval (optimize_nam c)).
  { unfold eval. rewrite Hnam.
    rewrite Mscale_assoc.
    rewrite Cexp_mul_neg_l.
    rewrite Mscale_1_l. auto. }
  rewrite Haux, Hmap.
  remember (swap_route (optimize_nam c) lay cg get_path) as c0.
  clear Hmap Hnam Haux.
  specialize (optimize_preserves_semantics (decompose_swaps c0 cg)) as Hopt.
  destruct Hopt as [y Hopt].
  apply decompose_swaps_preserves_WT.
  subst c0. apply swap_route_preserves_WT; auto.
  rewrite cast_same.
  apply optimize_nam_preserves_WT; auto.
  assert (Haux: eval (optimize (decompose_swaps c0 cg)) = Cexp y .* eval (decompose_swaps c0 cg)).
  { unfold eval. rewrite Hopt. reflexivity. }
  rewrite Haux.
  clear Hopt Haux.
  exists (- x - y)%R.
  distribute_scale.
  apply f_equal2.
  rewrite <- Cexp_add.
  field_simplify (- x - y + y)%R.
  reflexivity.
  apply f_equal2; try reflexivity.
  apply f_equal2; try reflexivity.
  symmetry.
  apply decompose_swaps_preserves_semantics.
  subst cg.
  reflexivity.
Qed.

Lemma optimize_and_map_to_lnn_ring_16_respects_constraints : forall (c : circ 16) c',
  optimize_and_map_to_lnn_ring_16 c = Some c' ->
  respects_constraints_directed (is_in_graph (make_lnn_ring 16)) U_CX c'.
Proof.
  intros c c' H.
  unfold optimize_and_map_to_lnn_ring_16 in H.
  remember (make_lnn_ring 16) as cg.
  remember (lnn_ring_path_finding_fun 16) as get_path.
  remember (greedy_layout c cg (lnn_ring_qubit_ordering_fun 16)) as lay.
  assert (Hpath : get_path_valid cg get_path).
  { subst. apply lnn_ring_path_finding_fun_valid. }
  clear Heqget_path.
  destruct (check_well_typed c 16) eqn:WT; inversion H.
  apply check_well_typed_correct in WT.
  replace (graph_dim cg) with 16 in *.
  rewrite cast_same in WT.
  assert (WF : layout_bijective 16 lay).
  { subst. apply greedy_layout_well_formed.
    rewrite cast_same. assumption.
    apply lnn_ring_qubit_ordering_fun_valid. }
  clear - WT Hpath WF Heqcg.
  apply optimize_preserves_mapping.
  apply decompose_swaps_respects_constraints.
  apply swap_route_respects_constraints_undirected; auto.
  replace (graph_dim cg) with 16.
  rewrite cast_same.
  apply optimize_nam_preserves_WT. auto.
  subst. reflexivity.
  replace (graph_dim cg) with 16. auto.
  subst. reflexivity.
  subst. reflexivity.
Qed.


Require Import VST.msl.base.
Require Import VST.msl.ageable.
Require Import VST.msl.functors.
Require Import VST.msl.predicates_hered.
Import VST.msl.functors.MixVariantFunctor.
Import VST.msl.functors.MixVariantFunctorLemmas.
Require Import Arith.

Module Type KNOT_INPUT__MIXVARIANT_HERED_T_OTH_REL.
  Parameter F : functor.

  Parameter other : Type.

  Parameter Rel : forall A, F A -> F A -> Prop.

  Parameter Rel_fmap : forall A B (f1: A->B) (f2:B->A) x y,
    Rel A x y ->
    Rel B (fmap F f1 f2 x) (fmap F f1 f2 y).
  Axiom Rel_refl : forall A x, Rel A x x.
  Axiom Rel_trans : forall A x y z,
    Rel A x y -> Rel A y z -> Rel A x z.

  Parameter ORel : other -> other -> Prop.
  Axiom ORel_refl : reflexive other ORel.
  Axiom ORel_trans : transitive other ORel.

  Parameter T:Type.
  Parameter T_bot:T.

  Parameter T_rel : T -> T -> Prop.
  Parameter T_rel_bot : forall x, T_rel T_bot x.
  Parameter T_rel_refl : forall x, T_rel x x.
  Parameter T_rel_trans : transitive T T_rel.

End KNOT_INPUT__MIXVARIANT_HERED_T_OTH_REL.

Module Type KNOT__MIXVARIANT_HERED_T_OTH_REL.
  Declare Module KI: KNOT_INPUT__MIXVARIANT_HERED_T_OTH_REL.
  Import KI.

  Parameter knot:Type.
  Parameter ageable_knot : ageable knot.
  #[global] Existing Instance ageable_knot.

  Parameter hered : (knot * other -> T) -> Prop.
  Definition predicate := { p:knot * other -> T | hered p }.

  Parameter squash : (nat * F predicate) -> knot.
  Parameter unsquash : knot -> (nat * F predicate).

  Parameter approx : nat -> predicate -> predicate.

  Axiom squash_unsquash : forall k:knot, squash (unsquash k) = k.
  Axiom unsquash_squash : forall (n:nat) (f:F predicate),
    unsquash (squash (n,f)) = (n, fmap F (approx n) (approx n) f).

  Axiom approx_spec : forall n p ko,
    proj1_sig (approx n p) ko =
     if (Compare_dec.le_gt_dec n (level (fst ko))) then T_bot else proj1_sig p ko.

  Definition knot_rel (k1 k2:knot) :=
    let (n,f) := unsquash k1 in
    let (n',f') := unsquash k2 in
    n = n' /\ Rel predicate f f'.

  Axiom knot_age1 : forall k:knot,
    age1 k =
    match unsquash k with
    | (O,_) => None
    | (S n,x) => Some (squash (n,x))
    end.

  Axiom knot_level : forall k:knot,
    level k = fst (unsquash k).

  Axiom hered_spec : forall p,
    hered p =
    (forall k k' k'' o o',
      clos_refl_trans _ age k k' ->
      knot_rel  k' k'' ->
      ORel o o' ->
      T_rel (p (k,o)) (p (k'',o'))).

  #[export] Program Instance ext_knot : Ext_ord knot := { ext_order := knot_rel }.
  Next Obligation.
  Proof.
    unfold knot_rel. split.
    - intros k.
      destruct (unsquash k); split; auto.
      apply Rel_refl.
    - intros k1 k2 k3.
      destruct (unsquash k1), (unsquash k2), (unsquash k3).
      intros [] []; subst; split; auto.
      eapply Rel_trans; eauto.
  Qed.
  Next Obligation.
  Proof.
    intros ?????.
    unfold age, knot_rel in *. rewrite knot_age1 in H.
    destruct (unsquash y) eqn: Hy.
    destruct n; inv H; simpl.
    destruct (unsquash z) eqn: Hz.
    destruct H0 as [? H0]; subst.
    exists (squash (n, _f0)).
    - rewrite !unsquash_squash.
      split; auto.
      apply Rel_fmap; auto.
    - rewrite knot_age1, Hz; auto.
  Qed.
  Next Obligation.
  Proof.
    unfold age, knot_rel in *. rewrite knot_age1 in H0.
    destruct (unsquash a) eqn: Ha.
    destruct n; inv H0.
    destruct (unsquash b) eqn: Hb.
    destruct H as [? H]; subst.
    exists (squash (n, _f0)).
    split.
    - rewrite knot_age1, Hb; auto.
    - rewrite !unsquash_squash; split; auto.
      apply Rel_fmap; auto.
  Qed.
  Next Obligation.
  Proof.
    rewrite !knot_level. unfold knot_rel in H.
    destruct (unsquash a), (unsquash b), H; auto.
  Qed.

  Lemma knot_order : ext_order = knot_rel.
  Proof. reflexivity. Qed.

End KNOT__MIXVARIANT_HERED_T_OTH_REL.

Module Knot_MixVariantHeredTOthRel (KI':KNOT_INPUT__MIXVARIANT_HERED_T_OTH_REL) :
  KNOT__MIXVARIANT_HERED_T_OTH_REL with Module KI:=KI'.
  Module KI := KI'.
  Import KI.

  Definition sinv_prod X := prod X (F X * other -> T).

  Definition guppy_sig := (fun X:Type => X * (F X * other -> T) -> Prop).
  Definition guppy_ty := sigT guppy_sig.

  Definition guppy_step_ty (Z:guppy_ty) : Type :=
    (sig (fun (x:sinv_prod (projT1 Z)) => projT2 Z x)).

  Definition guppy_age (Z:guppy_ty) (x:guppy_step_ty Z) : projT1 Z := fst (proj1_sig x).
  Definition guppy_unage (Z:guppy_ty)
    (H:forall t, projT2 Z (t,fun _ => T_bot))
    (x:projT1 Z) : guppy_step_ty Z :=
    exist (fun z => projT2 Z z) (x, fun _ => T_bot) (H x).

  Definition guppy_step_prop (Z:guppy_ty) (xf:sinv_prod (guppy_step_ty Z)) :=
    (forall (k:F (guppy_step_ty Z)) (o:other) H,
         T_rel (snd xf (k,o))
               (snd (proj1_sig (fst xf)) (fmap F (guppy_age Z) (guppy_unage Z H) k,o))) /\
    (forall (k k':F (guppy_step_ty Z)) (o o':other),
         Rel (guppy_step_ty Z) k k' ->
         ORel o o' ->
         T_rel (snd xf (k,o)) (snd xf (k',o'))).

  Definition guppy_step (Z:guppy_ty) : guppy_ty :=
    existT guppy_sig (guppy_step_ty Z) (guppy_step_prop Z).

  Definition guppy_base : guppy_ty :=
    existT guppy_sig unit
      (fun xf =>
        (forall (k k':F unit) (o o':other),
          Rel unit k k' ->
          ORel o o' ->
          T_rel (snd xf (k,o)) (snd xf (k',o')))).

  Fixpoint guppy (n:nat) : guppy_ty :=
    match n with
    | 0    => guppy_base
    | S n' => guppy_step (guppy n')
    end.

  Definition sinv (n:nat) : Type := projT1 (guppy n).
  Definition sinv_prop (n:nat) : prod (sinv n) (F (sinv n) * other -> T) -> Prop := projT2 (guppy n).

  Fixpoint floor (m:nat) (n:nat) (p:sinv (m+n)) : sinv n :=
    match m as m' return forall (p : sinv (m'+n)), sinv n with
    | O => fun p => p
    | S m' => fun p => floor m' n (fst (proj1_sig p))
    end p.

  Definition knot := { n:nat & F (sinv n) }.

  Definition sinv_age n : sinv (S n) -> sinv n := guppy_age (guppy n).
  Program Definition sinv_unage n : sinv n -> sinv (S n) := guppy_unage (guppy n) _.
  Next Obligation.
    revert t; induction n; simpl; auto.
    repeat intro.
    apply T_rel_bot.
    split; simpl in *; repeat intro.
    apply T_rel_bot.
    apply T_rel_bot.
  Qed.

  Definition F_sinv n := F (sinv n).

  Definition age1_def (k:knot) : option knot :=
    match k with
      | existT _ 0 f => None
      | existT _ (S m) f => Some
          (existT F_sinv m (fmap F (sinv_age m) (sinv_unage m) f))
    end.

  Definition age_def x y := age1_def x = Some y.

  Inductive knot_rel_inner : knot -> knot -> Prop :=
    | intro_krel : forall n (f f':F_sinv n),
         Rel _ f f' ->
         knot_rel_inner (existT (F_sinv) n f) (existT (F_sinv) n f').

  Definition hered (p:knot * other -> T) : Prop :=
    forall k k' k'' o o',
      clos_refl_trans _ age_def k k' ->
      knot_rel_inner k' k'' -> ORel o o' ->
      T_rel (p (k,o)) (p (k'',o')).

  Definition predicate := { p:knot * other -> T | hered p }.

  Definition app_sinv (n:nat) (p:sinv (S n)) (x:F_sinv n * other) :=
    snd (proj1_sig p) x.

  Section stratifies.
    Variable Q:knot * other -> T.
    Variable HQ:hered Q.

    Fixpoint stratifies (n:nat) : sinv n -> Prop :=
    match n as n' return sinv n' -> Prop with
    | 0 => fun _ => True
    | S n' => fun (p:sinv (S n')) =>
          stratifies n' (fst (proj1_sig p)) /\
          forall (k:F_sinv n') (o:other), snd (proj1_sig p) (k,o) = Q (existT F_sinv n' k,o)
    end.

    Lemma stratifies_unique : forall n p1 p2,
      stratifies n p1 ->
      stratifies n p2 ->
      p1 = p2.
    Proof.
      induction n; simpl; intuition.
      destruct p1; destruct p2; auto.
      destruct p1; destruct p2.
      simpl in *; fold guppy in *.
      cut (x = x0).
      intros.
      revert p p0 H2 H3.
      rewrite <- H0.
      intros.
      replace p0 with p by (apply proof_irr); auto.
      destruct x; destruct x0; simpl in *.
      apply injective_projections; simpl.
      apply IHn; auto.
      extensionality; intros.
      simpl in *.
      destruct x as [x o].
      destruct (H2 x o); destruct (H3 x o).
      rewrite H2.
      rewrite H3.
      auto.
    Qed.

    Definition stratify (n:nat) : { x:sinv n | stratifies n x }.
    Proof.
      induction n.
      exists tt; simpl; exact I.
      assert (HX:
        projT2 (guppy n)
        (proj1_sig IHn, fun v : F_sinv n * other => Q (existT F_sinv n (fst v),snd v))).
      destruct n.
      simpl; intros.
      eapply HQ.
      apply rt_refl.
      constructor; auto.
      auto.
      simpl; intros.
      destruct IHn; simpl.
      simpl in s; destruct s.
      destruct x; simpl in *; fold guppy in *.
      destruct x; simpl in *.
      split; hnf; simpl; intros.
      rewrite H0.
      eapply HQ.
      apply rt_step.
      hnf; simpl.
      reflexivity.
      constructor; auto.
      unfold sinv_unage.
      replace (sinv_unage_obligation_1 n) with H1.
      unfold sinv_age.
      apply Rel_refl.
      apply proof_irr.
      apply ORel_refl.
      eapply HQ.
      apply rt_refl.
      constructor; auto.
      auto.

      exists ((exist (fun x => projT2 (guppy n) x) ( proj1_sig IHn, fun v:F_sinv n * other => Q (existT (F_sinv) n (fst v),snd v) ) HX)).
      simpl; split; auto.
      destruct IHn; auto.
    Qed.
  End stratifies.

  Lemma decompose_nat : forall (x y:nat), { m:nat & y = (m + S x) } + { ge x y }.
  Proof.
    intros x y; revert x; induction y; simpl; intros.
    right; auto with arith.
    destruct (IHy x) as [[m H]|H].
    left; exists (S m); lia.
    destruct (Peano_dec.eq_nat_dec x y).
    left; exists O; lia.
    right; lia.
  Qed.

  Definition unstratify (n:nat) (p:sinv n) : knot * other -> T := fun w =>
    match w with (existT _ nw w',o) =>
      match decompose_nat nw n with
        | inleft (existT _ m Hm) => snd (proj1_sig (floor m (S nw) (eq_rect  n _ p (m + S nw) Hm))) (w',o)
        | inright H => T_bot
      end
    end.

  Lemma floor_shuffle:
    forall (m1 n : nat)
      (p1 : sinv (m1 + S n)) (H1 : (m1 + S n) = (S m1 + n)),
      floor (S m1) n (eq_rect (m1 + S n) sinv p1 (S m1 + n) H1) = fst (proj1_sig (floor m1 (S n) p1)).
  Proof.
    intros.
    remember (fst (proj1_sig (floor m1 (S n) p1))) as p.
    fold guppy in *.
    revert n p1 H1 p Heqp.
    induction m1; simpl; intros.
    replace H1 with (refl_equal (S n)) by (apply proof_irr); simpl; auto.
    assert (m1 + S n = S m1 + n) by lia.
    destruct p1 as [[p1 f'] Hp1]; simpl in *; fold guppy in *.
    generalize (IHm1 n p1 H p Heqp).
    clear.
    revert Hp1 H1; generalize H.
    revert p1 f'.
    rewrite H.
    simpl; intros.
    replace H1 with (refl_equal (S (S (m1 + n)))) by (apply proof_irr).
    simpl.
    replace H0 with (refl_equal (S (m1+n))) in H2 by (apply proof_irr).
    simpl in H2.
    trivial.
  Qed.

  Lemma unstratify_hered : forall n p,
    hered (unstratify n p).
  Proof.
    intros.
    hnf; intros.
    apply T_rel_trans with (unstratify n p (k',o)).
    clear o' H0 H1.
    induction H.
    hnf in H; simpl in H.
    destruct x as [x f]; simpl in H.
    destruct x; try discriminate.
    assert (y =
      (existT (F_sinv) x (fmap F (sinv_age x) (sinv_unage x) f))).
    inversion H; auto.
    subst y.
    unfold unstratify.
    case_eq (decompose_nat (S x) n); intros.
    destruct s.
    case_eq (decompose_nat x n); intros.
    destruct s.
    destruct n.
    exfalso; lia.
    assert (S x0 = x1) by lia; subst x1.
    revert H1.
    generalize e e0; revert p; rewrite e; intros.
    rewrite floor_shuffle.
    replace e1 with (refl_equal (x0 + S (S x)));
      simpl eq_rect.
    2: apply proof_irr.
    revert H1.
    generalize (floor x0 (S (S x)) p).
    intros [[s' fs] Hs] H1; simpl in *; fold guppy in *.
    destruct Hs.
    simpl in H2.
    eapply H2; auto.
    exfalso.
    lia.
    apply T_rel_bot.
    apply T_rel_refl.
    eapply T_rel_trans; eauto.

    clear H.
    inv H0.
    simpl.
    destruct (decompose_nat n0 n); [ | apply T_rel_bot ].
    destruct s; simpl.
    destruct (floor x (S n0) (eq_rect n sinv p (x +S n0) e)); simpl.
    destruct n0; simpl in x0; destruct x0; simpl.
    apply p0; auto.
    apply p0; auto.
  Qed.

  Lemma unstratify_Q : forall n (p:sinv n) Q,
    stratifies Q n p ->
    forall (k:knot) (o:other),
      projT1 k < n ->
      (unstratify n p (k,o) = Q (k,o)).
  Proof.
    intros.
    unfold unstratify.
    destruct k.
    destruct (decompose_nat x n).
    destruct s.
    simpl in H0.
    2: simpl in *; exfalso; lia.
    clear H0.
    revert p H.
    generalize e.
    rewrite e.
    intros.
    replace e0 with (refl_equal (x0 + S x)) by apply proof_irr.
    simpl.
    clear e e0.
    revert p H.
    induction x0; simpl; intros.
    destruct H.
    auto.
    destruct H.
    apply IHx0.
    auto.
  Qed.

  Lemma stratifies_unstratify_more :
    forall (n m1 m2:nat) (p1:sinv (m1+n)) (p2:sinv (m2+n)),
      floor m1 n p1 = floor m2 n p2 ->
      (stratifies (unstratify (m1+n) p1) n (floor m1 n p1) ->
       stratifies (unstratify (m2+n) p2) n (floor m2 n p2)).
  Proof.
    induction n; intuition.
    split.
    assert (m2 + S n = S m2 + n) by lia.
    erewrite <- floor_shuffle.
    instantiate (1:=H1).
    replace (unstratify (m2 + S n) p2)
      with (unstratify (S m2 + n) (eq_rect (m2 + S n) sinv p2 (S m2 + n) H1)).
    assert (m1 + S n = S m1 + n) by lia.
    eapply (IHn (S m1) (S m2)
      (eq_rect (m1 + S n) sinv p1 (S m1 + n) H2)).
    rewrite floor_shuffle.
    rewrite floor_shuffle.
    rewrite H; auto.
    clear - H0.
    rewrite floor_shuffle.
    simpl in H0.
    destruct H0.
    clear H0.
    revert p1 H.
    generalize H2.
    rewrite <- H2.
    intros.
    replace H0 with (refl_equal (m1 + S n)) by apply proof_irr; auto.
    clear.
    revert p2.
    generalize H1.
    rewrite H1.
    intros.
    replace H0 with (refl_equal (S m2 + n)) by apply proof_irr; auto.

    intros.
    simpl.
    destruct (decompose_nat n (m2 + S n)).
    destruct s.
    assert (m2 = x).
    lia.
    subst x.
    replace e with (refl_equal (m2 + S n)).
    simpl; tauto.
    apply proof_irr.
    exfalso; lia.
  Qed.

  Lemma stratify_unstratify : forall n p H,
    proj1_sig (stratify (unstratify n p) H n) = p.
  Proof.
    intros.
    apply stratifies_unique with (unstratify n p).
    destruct (stratify _ H n).
    simpl; auto.
    clear H.
    revert p; induction n.
    simpl; intros; auto.
    intros.
    simpl; split.

    assert (stratifies (unstratify n (fst (proj1_sig p))) n (fst (proj1_sig p))).
    apply IHn.
    apply (stratifies_unstratify_more n 0 1 (fst (proj1_sig p)) p).
    simpl; auto.
    auto.

    intros.
    destruct (decompose_nat n (S n)).
    destruct s.
    assert (x = 0) by lia.
    subst x.
    simpl.
    simpl in e.
    replace e with (refl_equal (S n)) by apply proof_irr.
    simpl.
    split; auto.
    exfalso; lia.
  Qed.

  Definition strat (n:nat) (p:predicate) : sinv n :=
    proj1_sig (stratify (proj1_sig p) (proj2_sig p) n).

  Definition unstrat (n:nat) (p:sinv n) : predicate :=
    exist hered (unstratify n p) (unstratify_hered n p).

  Definition squash (x:nat * F predicate) : knot :=
    match x with (n,f) => existT (F_sinv) n (fmap F (strat n) (unstrat n) f) end.

  Definition unsquash (k:knot) : nat * F predicate :=
    match k with existT _ n f => (n, fmap F (unstrat n) (strat n) f) end.

  Definition knot_level_def (k:knot) : nat :=
    fst (unsquash k).

  Definition knot_age1_def (k:knot) : option knot :=
    match unsquash k with
    | (O,_) => None
    | (S n,x) => Some (squash (n,x))
    end.

  Definition knot_unage_def (k:knot) :=
    let (n,k) := unsquash k in squash (S n,k).

  Program Definition approx (n:nat) (p:predicate) : predicate :=
    fun w => if (Compare_dec.le_gt_dec n (knot_level_def (fst w))) then T_bot else proj1_sig p w.
  Next Obligation.
    hnf; simpl; intros.
    destruct (Compare_dec.le_gt_dec n (knot_level_def k)).
    apply T_rel_bot.
    destruct (Compare_dec.le_gt_dec n (knot_level_def k'')).
    exfalso.
    cut (knot_level_def k'' <= knot_level_def k).
    lia.
    replace (knot_level_def k'') with (knot_level_def k').
    clear -H; induction H.
    hnf in H.
    unfold age1_def in H.
    destruct x; destruct y; simpl.
    destruct x; try discriminate.
    inv H.
    simpl.
    unfold knot_level_def; simpl; auto.
    auto.
    eapply Nat.le_trans; eauto.
    inv H0.
    unfold knot_level_def; simpl; auto.

    destruct p as [p Hp]; simpl.
    eapply Hp; eauto.
  Qed.

  Lemma strat_unstrat : forall n,
    strat n oo unstrat n = id (sinv n).
  Proof.
    intros; extensionality p.
    unfold compose, id.
    unfold strat, unstrat.
    simpl.
    rewrite stratify_unstratify.
    auto.
  Qed.

  Lemma predicate_eq : forall (p1 p2:predicate),
    proj1_sig p1 = proj1_sig p2 ->
    p1 = p2.
  Proof.
    intros; destruct p1; destruct p2; simpl in H.
    subst x0.
    replace h0 with h by apply proof_irr.
    auto.
  Qed.

  Lemma unstrat_strat : forall n,
    unstrat n oo strat n = approx n.
  Proof.
    intros.
    extensionality.
    unfold compose.
    unfold unstrat, strat.
    unfold approx.
    apply predicate_eq.
    simpl.
    extensionality k.
    destruct (Compare_dec.le_gt_dec n (knot_level_def (fst k))).
    unfold unstratify.
    destruct k.
    destruct k.
    unfold knot_level_def in l.
    simpl in *.
    destruct (decompose_nat x0 n); simpl.
    destruct s; simpl; exfalso; lia.
    auto.
    destruct x as [x Hx]; simpl.
    destruct (stratify x Hx n); simpl.
    destruct k.
    rewrite unstratify_Q with (Q:=x); auto.
    unfold level in *.
    destruct k; simpl in *; auto.
  Qed.

  Lemma squash_unsquash : forall k, squash (unsquash k) = k.
  Proof.
    intros.
    destruct k as [n f]; simpl.
    f_equal.
    change ((fmap F (strat n) (unstrat n) oo (fmap F (unstrat n) (strat n))) f = f).
    rewrite fmap_comp.
    rewrite strat_unstrat.
    rewrite fmap_id.
    auto.
  Qed.

  Lemma unsquash_squash : forall n f,
    unsquash (squash (n,f)) = (n, fmap F (approx n) (approx n) f).
  Proof.
    intros.
    unfold unsquash, squash.
    f_equal.
    change ((fmap F (unstrat n) (strat n) oo (fmap F (strat n) (unstrat n))) f = fmap F (approx n) (approx n) f).
    rewrite fmap_comp.
    rewrite unstrat_strat.
    auto.
  Qed.

  Lemma strat_Sx_unstrat : forall x,
    sinv_unage x = strat (S x) oo unstrat x.
  Proof.
    intros.
    extensionality k.
    unfold sinv_unage.
    generalize (sinv_unage_obligation_1); intro P.
    unfold guppy_unage.
    unfold compose, strat, unstrat.
    simpl.
    apply stratifies_unique with (unstratify x k).
    revert k.
    induction x; simpl; intuition.
    destruct (decompose_nat 0 0); auto.
    destruct s; exfalso; lia.
    eapply (stratifies_unstratify_more x 0 1).
    simpl; reflexivity.
    simpl.
    simpl in *.
    destruct (IHx (fst (proj1_sig k))); auto.
    destruct (decompose_nat x (S x)).
    destruct s.
    assert (x0 = 0) by lia; subst x0.
    simpl in *.
    replace e with (refl_equal (S x)) by apply proof_irr; auto.
    exfalso; lia.
    destruct (decompose_nat x (S x)).
    destruct s.
    assert (x0 = 0) by lia; subst x0.
    simpl in *.
    destruct (decompose_nat (S x) (S x)).
    destruct s; exfalso; lia.
    auto.
    destruct (decompose_nat (S x) (S x)).
    destruct s; exfalso; lia.
    auto.

    destruct (stratify (unstratify x k) (unstratify_hered x k) (S x)).
    simpl stratifies in s; case s; intros.
    simpl stratifies; split; auto.
  Qed.

  Lemma strat_unstrat_Sx : forall x,
    sinv_age x = strat x oo unstrat (S x).
  Proof.
    intros.
    extensionality k.
    unfold sinv_age, guppy_age.
    unfold compose.
    unfold strat, unstrat.
    simpl.
    apply stratifies_unique with (unstratify x (fst (proj1_sig k))).
    revert k; induction x; simpl; auto.
    intros.
    split.
    eapply (stratifies_unstratify_more x 0 1 ).
    simpl; reflexivity.
    simpl.
    apply IHx.
    intros.
    destruct (decompose_nat x (S x)).
    destruct s.
    assert (x0 = 0) by lia; subst x0.
    simpl in *.
    replace e with (refl_equal (S x)) by apply proof_irr; simpl.
    tauto.
    exfalso; lia.
    destruct (stratify (unstratify (S x) k)
      (unstratify_hered (S x) k) x).
    simpl; auto.
    cut (x0 = (fst (proj1_sig k))); intros.
    subst x0.
    eapply (stratifies_unstratify_more x 1 0).
    simpl; reflexivity.
    simpl; auto.
    eapply stratifies_unique.
    apply s.
    eapply (stratifies_unstratify_more x 0 1).
    simpl; reflexivity.
    simpl.
    generalize (fst (proj1_sig k) : sinv x).
    clear.
    induction x; simpl; intuition.
    eapply (stratifies_unstratify_more x 0 1).
    simpl; reflexivity.
    simpl.
    apply IHx.
    destruct (decompose_nat x (S x)).
    destruct s0.
    assert (x0 = 0) by lia; subst.
    simpl in *.
    replace e with (refl_equal (S x)); simpl; auto.
    apply proof_irr.
    exfalso; lia.
  Qed.

  Lemma age1_eq : forall k,
    age1_def k = knot_age1_def k.
  Proof.
    intros.
    unfold knot_age1_def.
    case_eq (unsquash k); intros.
    case_eq k; intros.
    simpl.
    assert (n = x).
    subst k.
    inv H; auto.
    subst x.
    destruct n; auto.
    f_equal.
    f_equal.

    rewrite strat_Sx_unstrat.
    rewrite strat_unstrat_Sx.
    rewrite <- fmap_comp.
    unfold compose.
    f_equal.
    subst k.
    inv H; auto.
  Qed.

  Lemma unsquash_inj : forall k1 k2,
    unsquash k1 = unsquash k2 ->
    k1 = k2.
  Proof.
    intros.
    rewrite <- (squash_unsquash k1).
    rewrite <- (squash_unsquash k2).
    rewrite H.
    trivial.
  Qed.
  Arguments unsquash_inj [k1 k2] _.

  Lemma pred_ext : forall (p1 p2:predicate),
    (forall x, proj1_sig p1 x = proj1_sig p2 x) ->
    p1 = p2.
  Proof.
    intros.
    destruct p1 as [p1 Hp1]; destruct p2 as [p2 Hp2].
    simpl in *.
    assert (p1 = p2).
    extensionality x; auto.
    subst p2.
    replace Hp2 with Hp1; auto.
    apply proof_irr.
  Qed.

  Lemma approx_spec : forall n p ko,
    proj1_sig (approx n p) ko =
     if (Compare_dec.le_gt_dec n (knot_level_def (fst ko))) then T_bot else proj1_sig p ko.
  Proof.
    intros; simpl; auto.
  Qed.

  Lemma ag_knot_facts : ageable_facts knot knot_level_def knot_age1_def.
  Proof.
    constructor.

    unfold knot_age1_def; unfold knot_level_def; simpl; intros x'.
    destruct (unsquash x') as [n f] eqn:?H; intros.
    destruct x' as [x f0].
    exists (squash (S x, fmap F (unstrat x) (strat x) f0)).
    rewrite unsquash_squash.
    f_equal. f_equal.
    clear.
    transitivity ((fmap F (strat x) (unstrat x) oo fmap F (approx (S x)) (approx (S x)) oo fmap F (unstrat x) (strat x)) f0); auto.
    do 2 rewrite fmap_comp.
    rewrite compose_assoc.
    replace (strat x oo approx (S x) oo unstrat x) with (@id (sinv x)).
    rewrite fmap_id. auto.
    rewrite <- (strat_unstrat x).
    f_equal.
    extensionality a.
    unfold compose, approx.
    case_eq (unstrat x a); intros.
    match goal with
      [ |- _ = exist _ ?X _ ] =>
      assert (x0 = X)
    end.
   2:{
    generalize (approx_obligation_1 (S x)
      (exist (fun p => hered p) x0 h)).
    rewrite <- H0.
    intros. f_equal.
    }
    extensionality.
    destruct x1.
    unfold unstrat in H.
    inv H.
    destruct k.
    unfold unstratify.
    unfold knot_level_def.
    simpl fst.
    destruct (decompose_nat x0 x).
    destruct s.
    destruct (Compare_dec.le_gt_dec (S x) x0).
    exfalso; lia.
    simpl.
    destruct (decompose_nat x0 x).
    destruct s.
    assert (x1 = x2) by lia.
    subst x2.
    replace e0 with e by apply proof_irr.
    auto.
    exfalso; lia.
    destruct (Compare_dec.le_gt_dec (S x) x0); auto.
    simpl.
    destruct (decompose_nat x0 x); auto.
    destruct s. exfalso. lia.

    intro.
    unfold knot_age1_def, knot_level_def.
    case_eq (unsquash x); intros.
    destruct n; simpl; intuition;
      discriminate.

    intros.
    unfold knot_age1_def, knot_level_def in *.
    case_eq (unsquash x); intros; rewrite H0 in H.
    destruct n; try discriminate; simpl.
    inv H; simpl; auto.
  Qed.

  Definition ageable_knot : ageable knot :=
    mkAgeable knot knot_level_def knot_age1_def ag_knot_facts.
  #[global] Existing Instance ageable_knot.

  Definition knot_rel (k1 k2:knot) :=
    let (n,f) := unsquash k1 in
    let (n',f') := unsquash k2 in
    n = n' /\ Rel predicate f f'.

  Lemma hered_spec : forall p,
    hered p =
    (forall k k' k'' o o',
      clos_refl_trans _ age k k' ->
      knot_rel  k' k'' ->
      ORel o o' ->
      T_rel (p (k,o)) (p (k'',o'))).
  Proof.
    intros.
    apply prop_ext.
    intuition.
    eapply H.
    instantiate (1:=k').
    clear -H0; induction H0; auto.
    apply rt_step.
    unfold age_def.
    rewrite age1_eq.
    auto.
    eapply rt_trans; eauto.
    destruct k' as [x f], k'' as [x0 f0].
    unfold knot_rel, unsquash in H1.
    destruct H1; subst.
    constructor.
    apply (Rel_fmap _ _ (strat x0) (unstrat x0)) in H3.
    change f with (id _ f).
    change f0 with (id _ f0).
    rewrite <- fmap_id.
    rewrite <- (strat_unstrat x0).
    rewrite <- fmap_comp.
    auto.
    assumption.

    hnf; intros.
    apply (H k k' k''); auto.
    clear -H0; induction H0; auto.
    apply rt_step.
    hnf.
    rewrite <- age1_eq; auto.
    eapply rt_trans; eauto.

    destruct k'; destruct k''.
    inv H1.
    simpl.
    hnf; split; auto.
    apply Eqdep_dec.inj_pair2_eq_dec in H5; auto.
    apply Eqdep_dec.inj_pair2_eq_dec in H7; auto.
    subst.
    apply Rel_fmap; auto.
    exact Peano_dec.eq_nat_dec.
    exact Peano_dec.eq_nat_dec.
  Qed.

  Lemma knot_age1 : forall k:knot,
    age1 k =
    match unsquash k with
    | (O,_) => None
    | (S n,x) => Some (squash (n,x))
    end.
  Proof.
    intros; reflexivity.
  Qed.

  Lemma knot_level : forall k:knot,
    level k = fst (unsquash k).
  Proof.
    intros; reflexivity.
  Qed.

  #[export] Program Instance ext_knot : Ext_ord knot := { ext_order := knot_rel }.
  Next Obligation.
  Proof.
    unfold knot_rel. split.
    - intros k.
      destruct (unsquash k); split; auto.
      apply Rel_refl.
    - intros k1 k2 k3.
      destruct (unsquash k1), (unsquash k2), (unsquash k3).
      intros [] []; subst; split; auto.
      eapply Rel_trans; eauto.
  Qed.
(*  Next Obligation.
  Proof.
    intros ?????.
    unfold age, knot_rel in *. rewrite knot_age1 in H0.
    destruct (unsquash z) eqn: Hz.
    destruct n; inv H0.
    rewrite unsquash_squash in H.
    destruct (unsquash x) eqn: Hx.
    destruct H as [? H]; subst.
    exists (squash (S n0, _f0)); simpl.
    - rewrite knot_age1, unsquash_squash.
      f_equal; apply unsquash_inj.
      rewrite unsquash_squash, Hx.
      rewrite fmap_app, <- (approx_approx1 1), <- (approx_approx2 1), <- (unsquash_approx Hx).
      reflexivity.
    - rewrite unsquash_squash; split; auto.
      rewrite (unsquash_approx Hz); rewrite (unsquash_approx Hx) in *.
      rewrite fmap_app, <- (approx_approx1 1), <- (approx_approx2 1).
      (* may not be true: the unaged pred may not be in Rel even if the aged one is *)
      admit.
  Admitted.*)
  Next Obligation.
  Proof.
    intros ?????.
    unfold age, knot_rel in *. rewrite knot_age1 in H.
    destruct (unsquash y) eqn: Hy.
    destruct n; inv H; simpl.
    destruct (unsquash z) eqn: Hz.
    destruct H0 as [? H0]; subst.
    exists (squash (n, _f0)); simpl.
    - split; auto.
      do 2 apply Rel_fmap; auto.
    - rewrite knot_age1, Hz; auto.
  Qed.
  Next Obligation.
  Proof.
    unfold age, knot_rel in *. rewrite knot_age1 in H0.
    destruct (unsquash a) eqn: Ha.
    destruct n; inv H0.
    destruct (unsquash b) eqn: Hb.
    destruct H as [? H]; subst.
    exists (squash (n, _f0)).
    split.
    - rewrite knot_age1, Hb; auto.
    - rewrite !unsquash_squash; split; auto.
      rewrite fmap_app, unstrat_strat.
      apply Rel_fmap; auto.
  Qed.
  Next Obligation.
  Proof.
    rewrite !knot_level. unfold knot_rel in H.
    destruct (unsquash a), (unsquash b), H; auto.
  Qed.

  Lemma knot_order : ext_order = knot_rel.
  Proof. reflexivity. Qed.

End Knot_MixVariantHeredTOthRel.

Module KnotLemmas1.

Class Input: Type := {
  knot: Type;
  Fpred: Type;
  squash: nat * Fpred -> knot;
  unsquash: knot -> nat * Fpred;
  approxF: nat -> Fpred -> Fpred;
  squash_unsquash : forall k:knot, squash (unsquash k) = k;
  unsquash_squash : forall (n:nat) (f:Fpred),
    unsquash (squash (n,f)) = (n, approxF n f)
}.

Class Output (input: Input): Prop := {
  unsquash_inj : forall k1 k2,
    unsquash k1 = unsquash k2 ->
    k1 = k2;
  squash_surj : forall k, exists n, exists Fp,
    squash (n, Fp) = k;
  unsquash_approx : forall k n Fp,
    unsquash k = (n, Fp) ->
    Fp = approxF n Fp
}.

Lemma Proof (kli: Input): Output kli.
Proof.
  constructor.
  + intros.
    rewrite <- (squash_unsquash k1).
    rewrite <- (squash_unsquash k2).
    rewrite H.
    trivial.
  + intros.
    remember (unsquash k).
    destruct p as [n f].
    exists n.
    exists f.
    rewrite Heqp.
    rewrite squash_unsquash.
    trivial.
  + intros.
    generalize H; intro.
    rewrite <- (squash_unsquash k) in H.
    rewrite H0 in H.
    rewrite unsquash_squash in H.
    inversion H.
    rewrite H2.
    symmetry.
    trivial.
Qed.

End KnotLemmas1.

Module KnotLemmas2.

Class Input: Type := {
  knot: Type;
  other: Type;
  T: Type;
  t0: T;
  ageable_knot : ageable knot;
  predicate: Type;
  p2p: predicate -> (knot * other -> T);
  approx : nat -> predicate -> predicate;
  pred_ext : forall (p1 p2:predicate),
    (forall x, p2p p1 x = p2p p2 x) ->
    p1 = p2;
  approx_spec : forall n p ko,
    p2p (approx n p) ko =
     if (Compare_dec.le_gt_dec n (level (fst ko))) then t0 else p2p p ko
}.

Class Output (input: Input): Prop := {
  approx_approx1 : forall m n,
    approx n = approx n oo approx (m+n);
  approx_approx2 : forall m n,
    approx n = approx (m+n) oo approx n
}.

Lemma Proof (kli: Input): Output kli.
Proof.
  constructor.
  + intros.
    extensionality p.
    apply pred_ext.
    intros [k o].
    unfold compose.
    repeat rewrite approx_spec.
    simpl.
    destruct (Compare_dec.le_gt_dec n (level k)); auto.
    destruct (Compare_dec.le_gt_dec (m+n) (level k)); auto.
    exfalso; lia.
  + intros.
    extensionality p.
    apply pred_ext.
    intros [k o].
    unfold compose.
    repeat rewrite approx_spec.
    simpl.
    destruct (Compare_dec.le_gt_dec (m+n) (level k)); auto.
    destruct (Compare_dec.le_gt_dec n (level k)); auto.
    exfalso; lia.
Qed.

End KnotLemmas2.

Module KnotLemmas_MixVariantHeredTOthRel (K : KNOT__MIXVARIANT_HERED_T_OTH_REL).
  Import K.KI.
  Import K.

  Lemma unsquash_inj : forall k1 k2,
    unsquash k1 = unsquash k2 ->
    k1 = k2.
  Proof.
    apply
     (@KnotLemmas1.unsquash_inj
       (KnotLemmas1.Build_Input _ _ _ _ _ squash_unsquash unsquash_squash)),
     (KnotLemmas1.Proof).
  Qed.
  Arguments unsquash_inj [k1 k2] _.

  Lemma squash_surj : forall k, exists n, exists Fp,
    squash (n, Fp) = k.
  Proof.
    apply
     (@KnotLemmas1.squash_surj
       (KnotLemmas1.Build_Input _ _ _ _ _ squash_unsquash unsquash_squash)),
     (KnotLemmas1.Proof).
  Qed.

  Lemma unsquash_approx : forall k n Fp,
    unsquash k = (n, Fp) ->
    Fp = fmap F (approx n) (approx n) Fp.
  Proof.
    apply
     (@KnotLemmas1.unsquash_approx
       (KnotLemmas1.Build_Input _ _ _ _ _ squash_unsquash unsquash_squash)),
     (KnotLemmas1.Proof).
  Qed.
  Arguments unsquash_approx [k n Fp] _.

  Lemma pred_ext : forall (p1 p2:predicate),
    (forall x, proj1_sig p1 x = proj1_sig p2 x) ->
    p1 = p2.
  Proof.
    intros.
    destruct p1 as [p1 Hp1]; destruct p2 as [p2 Hp2].
    simpl in *.
    assert (p1 = p2).
    extensionality x; auto.
    subst p2.
    replace Hp2 with Hp1; auto.
    apply proof_irr.
  Qed.

  Lemma approx_approx1 : forall m n,
    approx n = approx n oo approx (m+n).
  Proof.
    apply
     (@KnotLemmas2.approx_approx1
       (KnotLemmas2.Build_Input _ _ _ _ _ _ _ _ pred_ext approx_spec)),
     (KnotLemmas2.Proof).
  Qed.

  Lemma approx_approx2 : forall m n,
    approx n = approx (m+n) oo approx n.
  Proof.
    apply
     (@KnotLemmas2.approx_approx2
       (KnotLemmas2.Build_Input _ _ _ _ _ _ _ _ pred_ext approx_spec)),
     (KnotLemmas2.Proof).
  Qed.

End KnotLemmas_MixVariantHeredTOthRel.

Module Type KNOT_FULL_OUTPUT.
  Declare Module KI: KNOT_INPUT__MIXVARIANT_HERED_T_OTH_REL.
  Declare Module K0: KNOT__MIXVARIANT_HERED_T_OTH_REL with Module KI := KI.
  Import K0.
  Parameter predicate: Type.
  Parameter pkp: bijection predicate K0.predicate.
End KNOT_FULL_OUTPUT.

Module Type KNOT_FULL.
  Declare Module KI: KNOT_INPUT__MIXVARIANT_HERED_T_OTH_REL.
  Declare Module KO: KNOT_FULL_OUTPUT with Module KI := KI.
  Import KI.
  Import KO.

  Definition knot : Type := KO.K0.knot.
  Definition ageable_knot : ageable knot := KO.K0.ageable_knot.
  #[global] Existing Instance ageable_knot.
  Definition ext_knot : Ext_ord knot := KO.K0.ext_knot.
  #[global] Existing Instance ext_knot.
  Definition predicate: Type := KO.predicate.

  Definition squash : (nat * KI.F predicate) -> knot :=
    fun k => KO.K0.squash
     (fst k, fmap KI.F (bij_f _ _ KO.pkp) (bij_g _ _ KO.pkp) (snd k)).

  Definition unsquash : knot -> (nat * KI.F predicate) :=
    fun k => let (n, f) := KO.K0.unsquash k in
      (n, fmap KI.F (bij_g _ _ KO.pkp) (bij_f _ _ KO.pkp) f).

  Parameter approx : nat -> predicate -> predicate.

  Axiom squash_unsquash : forall k:knot, squash (unsquash k) = k.
  Axiom unsquash_squash : forall (n:nat) (f:F predicate),
    unsquash (squash (n,f)) = (n, fmap F (approx n) (approx n) f).

  Axiom approx_spec : forall n p ko,
    proj1_sig (bij_f _ _ KO.pkp (approx n p)) ko =
     if (Compare_dec.le_gt_dec n (level (fst ko)))
     then KI.T_bot
     else proj1_sig (bij_f _ _ KO.pkp p) ko.

  Axiom knot_age1 : forall k:knot,
    age1 k =
    match unsquash k with
    | (O,_) => None
    | (S n,x) => Some (squash (n,x))
    end.

  Axiom knot_level : forall k:knot,
    level k = fst (unsquash k).

  Definition knot_rel (k1 k2:knot) :=
    let (n,f) := unsquash k1 in
    let (n',f') := unsquash k2 in
    n = n' /\ KI.Rel predicate f f'.

  Axiom knot_rel_spec: forall k1 k2: knot,
    knot_rel k1 k2 = KO.K0.knot_rel k1 k2.

End KNOT_FULL.

Module Type KNOT_FULL_LEMMAS.
  Declare Module K: KNOT_FULL.
  Import K.

  Axiom unsquash_inj : forall k1 k2,
    unsquash k1 = unsquash k2 ->
    k1 = k2.
  Arguments unsquash_inj [k1 k2] _.

  Axiom squash_surj : forall k, exists n, exists Fp,
    squash (n, Fp) = k.

  Axiom unsquash_approx : forall k n Fp,
    unsquash k = (n, Fp) ->
    Fp = fmap KI.F (approx n) (approx n) Fp.
  Arguments unsquash_approx [k n Fp] _.

  Axiom approx_approx1 : forall m n,
    approx n = approx n oo approx (m+n).

  Axiom approx_approx2 : forall m n,
    approx n = approx (m+n) oo approx n.

End KNOT_FULL_LEMMAS.

Module KnotFull
  (KI': KNOT_INPUT__MIXVARIANT_HERED_T_OTH_REL)
  (KO': KNOT_FULL_OUTPUT with Module KI := KI'):
  KNOT_FULL with Module KI := KI' with Module KO:=KO'.

  Import MixVariantFunctor.
  Module KI:=KI'.
  Module KO:=KO'.

  Definition knot: Type := KO.K0.knot.
  Definition ageable_knot : ageable knot := KO.K0.ageable_knot.
  #[global] Existing Instance ageable_knot.
  Definition ext_knot : Ext_ord knot := KO.K0.ext_knot.
  #[global] Existing Instance ext_knot.
  Definition predicate: Type := KO.predicate.

  Definition squash : (nat * KI.F predicate) -> knot :=
    fun k => KO.K0.squash
     (fst k, fmap KI.F (bij_f _ _ KO.pkp) (bij_g _ _ KO.pkp) (snd k)).

  Definition unsquash : knot -> (nat * KI.F predicate) :=
    fun k => let (n, f) := KO.K0.unsquash k in
      (n, fmap KI.F (bij_g _ _ KO.pkp) (bij_f _ _ KO.pkp) f).

  Definition approx : nat -> predicate -> predicate :=
    fun n => (bij_g _ _ KO.pkp) oo KO.K0.approx n oo (bij_f _ _ KO.pkp).

  Lemma squash_unsquash : forall k:knot, squash (unsquash k) = k.
  Proof.
    intros; unfold squash, unsquash.
    destruct (KO.K0.unsquash k) as [n f] eqn:?H; simpl.
    rewrite fmap_app, bij_fg_id, fmap_id.
    unfold id.
    rewrite <- H; apply KO.K0.squash_unsquash.
  Qed.

  Lemma unsquash_squash : forall (n:nat) (f:KI.F predicate),
    unsquash (squash (n,f)) = (n, fmap KI.F (approx n) (approx n) f).
  Proof.
    intros; unfold squash, unsquash, approx; simpl.
    rewrite KO.K0.unsquash_squash, !fmap_app, compose_assoc.
    auto.
  Qed.

  Lemma approx_spec : forall n p ko,
    proj1_sig (bij_f _ _ KO.pkp (approx n p)) ko =
     if (Compare_dec.le_gt_dec n (level (fst ko)))
     then KI.T_bot
     else proj1_sig (bij_f _ _ KO.pkp p) ko.
  Proof.
    intros.
    rewrite <- KO.K0.approx_spec.
    unfold approx.
    pattern (KO.K0.approx n) at 2.
    rewrite <- (id_unit2 _ _ (KO.K0.approx n)), <- (bij_fg_id KO.pkp).
    reflexivity.
  Qed.

  Lemma knot_age1 : forall k:knot,
    age1 k =
    match unsquash k with
    | (O,_) => None
    | (S n,x) => Some (squash (n,x))
    end.
  Proof.
    intros.
    unfold squash, unsquash.
    rewrite KO.K0.knot_age1.
    destruct (KO.K0.unsquash k) as [n f] eqn:?H.
    destruct n; auto.
    f_equal; simpl.
    rewrite fmap_app, bij_fg_id, fmap_id.
    auto.
  Qed.

  Lemma knot_level: forall k:knot, level k = fst (unsquash k).
  Proof.
    intros.
    unfold unsquash.
    rewrite KO.K0.knot_level.
    destruct (KO.K0.unsquash k) as [n f]; auto.
  Qed.

  Definition knot_rel (k1 k2:knot) :=
    let (n,f) := unsquash k1 in
    let (n',f') := unsquash k2 in
    n = n' /\ KI.Rel predicate f f'.

  Lemma knot_rel_spec: forall k1 k2: knot,
    knot_rel k1 k2 = KO.K0.knot_rel k1 k2.
  Proof.
    intros.
    unfold knot_rel, KO.K0.knot_rel, unsquash.
    destruct (KO.K0.unsquash k1) as [n1 f1].
    destruct (KO.K0.unsquash k2) as [n2 f2].
    f_equal.
    apply prop_ext.
    split; intros.
    + pose proof KI.Rel_fmap _ _ (bij_f _ _ KO.pkp) (bij_g _ _ KO.pkp) _ _ H.
      rewrite !fmap_app, bij_fg_id, fmap_id in H0.
      auto.
    + pose proof KI.Rel_fmap _ _ (bij_g _ _ KO.pkp) (bij_f _ _ KO.pkp) _ _ H.
      auto.
  Qed.

End KnotFull.

Module KnotFullLemmas (K: KNOT_FULL).
  Import K.KI.
  Import K.

  Lemma unsquash_inj : forall k1 k2,
    unsquash k1 = unsquash k2 ->
    k1 = k2.
  Proof.
    apply
     (@KnotLemmas1.unsquash_inj
       (KnotLemmas1.Build_Input _ _ _ _ _ squash_unsquash unsquash_squash)),
     (KnotLemmas1.Proof).
  Qed.
  Arguments unsquash_inj [k1 k2] _.

  Lemma squash_surj : forall k, exists n, exists Fp,
    squash (n, Fp) = k.
  Proof.
    apply
     (@KnotLemmas1.squash_surj
       (KnotLemmas1.Build_Input _ _ _ _ _ squash_unsquash unsquash_squash)),
     (KnotLemmas1.Proof).
  Qed.

  Lemma unsquash_approx : forall k n Fp,
    unsquash k = (n, Fp) ->
    Fp = fmap F (approx n) (approx n) Fp.
  Proof.
    apply
     (@KnotLemmas1.unsquash_approx
       (KnotLemmas1.Build_Input _ _ _ _ _ squash_unsquash unsquash_squash)),
     (KnotLemmas1.Proof).
  Qed.
  Arguments unsquash_approx [k n Fp] _.

  Lemma pred_ext : forall (p1 p2:predicate),
    (forall x, proj1_sig (bij_f _ _ KO.pkp p1) x =
               proj1_sig (bij_f _ _ KO.pkp p2) x) ->
    p1 = p2.
  Proof.
    intros.
    change (p1 = p2) with (id K.KO.predicate p1 = id K.KO.predicate p2).
    rewrite <- (bij_gf_id KO.pkp); unfold compose.
    destruct (bij_f _ _ KO.pkp p1) as [pp1 Hp1];
    destruct (bij_f _ _ KO.pkp p2) as [pp2 Hp2].
    simpl in *.
    assert (pp1 = pp2).
    extensionality x; auto.
    subst pp2.
    replace Hp2 with Hp1; auto.
    apply proof_irr.
  Qed.

  Lemma approx_approx1 : forall m n,
    approx n = approx n oo approx (m+n).
  Proof.
    apply
     (@KnotLemmas2.approx_approx1
       (KnotLemmas2.Build_Input _ _ _ _ _ _
         (@proj1_sig _ _ oo bij_f _ _ K.KO.pkp) _ pred_ext approx_spec)),
     (KnotLemmas2.Proof).
  Qed.

  Lemma approx_approx2 : forall m n,
    approx n = approx (m+n) oo approx n.
  Proof.
    apply
     (@KnotLemmas2.approx_approx2
       (KnotLemmas2.Build_Input _ _ _ _ _ _
         (@proj1_sig _ _ oo bij_f _ _ K.KO.pkp) _ pred_ext approx_spec)),
     (KnotLemmas2.Proof).
  Qed.

End KnotFullLemmas.

(*
Module Type KNOT_FULL_INPUT.
  Parameter F : functor.

  Parameter other : Type.

  Parameter Rel : forall A, F A -> F A -> Prop.

  Parameter Rel_fmap : forall A B (f1: A->B) (f2:B->A) x y,
    Rel A x y ->
    Rel B (fmap F f1 f2 x) (fmap F f1 f2 y).
  Axiom Rel_refl : forall A x, Rel A x x.
  Axiom Rel_trans : forall A x y z,
    Rel A x y -> Rel A y z -> Rel A x z.

  Parameter ORel : other -> other -> Prop.
  Axiom ORel_refl : reflexive other ORel.
  Axiom ORel_trans : transitive other ORel.

  Parameter T:Type.
  Parameter T_bot:T.

  Parameter T_rel : T -> T -> Prop.
  Parameter T_rel_bot : forall x, T_rel T_bot x.
  Parameter T_rel_refl : forall x, T_rel x x.
  Parameter T_rel_trans : transitive T T_rel.

  Parameter Pred: forall K: Type, ageable K -> (K -> K -> Prop) -> Type.

  Parameter Pred2predicate: forall {K agK KRel},
    Pred K agK KRel ->
    { p: K * other -> T |
      (forall k k' k'' o o',
      clos_refl_trans _ age k k' ->
      KRel k' k'' ->
      ORel o o' ->
      T_rel (p (k,o)) (p (k'',o'))) }.

  Parameter predicate2Pred: forall {K agK} {KRel: K -> K -> Prop},
    { p: K * other -> T |
      (forall (k k' k'': K) o o',
      clos_refl_trans _ age k k' ->
      KRel k' k'' ->
      ORel o o' ->
      T_rel (p (k,o)) (p (k'',o'))) } ->
    Pred K agK KRel.

  Axiom P2p2P: forall K agK KRel (P: Pred K agK KRel),
    predicate2Pred (Pred2predicate P) = P.

  Axiom p2P2p: forall K agK KRel p,
    Pred2predicate (@predicate2Pred K agK KRel p) = p.

End KNOT_FULL_INPUT.

Module Type KNOT_FULL.
  Declare Module KI: KNOT_FULL_INPUT.
  Import KI.

  Parameter knot:Type.
  Parameter ageable_knot : ageable knot.
  #[global] Existing Instance ageable_knot.
  Parameter knot_rel: knot -> knot -> Prop.

  Definition predicate: Type := Pred knot ageable_knot knot_rel.

  Parameter squash : (nat * F predicate) -> knot.
  Parameter unsquash : knot -> (nat * F predicate).

  Parameter approx : nat -> predicate -> predicate.

  Axiom squash_unsquash : forall k:knot, squash (unsquash k) = k.
  Axiom unsquash_squash : forall (n:nat) (f:F predicate),
    unsquash (squash (n,f)) = (n, fmap F (approx n) (approx n) f).

  Axiom approx_spec : forall n p ko,
    proj1_sig (Pred2predicate (approx n (predicate2Pred p))) ko =
     if (le_gt_dec n (level (fst ko))) then T_bot else proj1_sig p ko.

  Axiom knot_rel_spec: forall (k1 k2:knot),
     knot_rel k1 k2 =
    let (n,f) := unsquash k1 in
    let (n',f') := unsquash k2 in
    n = n' /\ Rel predicate f f'.

  Axiom knot_age1 : forall k:knot,
    age1 k =
    match unsquash k with
    | (O,_) => None
    | (S n,x) => Some (squash (n,x))
    end.

  Axiom knot_level : forall k:knot,
    level k = fst (unsquash k).

End KNOT_FULL.

Module KnotFull (KI': KNOT_FULL_INPUT): KNOT_FULL with Module KI:=KI'.

  Import MixVariantFunctor.
  Module KI:=KI'.

  Module Input.
    Definition F: functor := KI.F.
    Definition other: Type := KI.other.
    Definition Rel: forall (A:Type), KI.F A -> KI.F A -> Prop := KI.Rel.

    Definition Rel_fmap : forall A B (f1: A->B) (f2:B->A) x y,
      Rel A x y ->
      Rel B (fmap F f1 f2 x) (fmap F f1 f2 y)
      := KI.Rel_fmap.

    Definition Rel_refl : forall A x, Rel A x x := KI.Rel_refl.

    Definition Rel_trans : forall A x y z,
      Rel A x y -> Rel A y z -> Rel A x z
      := KI.Rel_trans.

    Definition ORel := KI.ORel.
    Definition ORel_refl := KI.ORel_refl.
    Definition ORel_trans := KI.ORel_trans.

    Definition T := KI.T.
    Definition T_bot: T := KI.T_bot.
    Definition T_rel : T -> T -> Prop := KI.T_rel.
    Definition T_rel_bot : forall x, T_rel T_bot x := KI.T_rel_bot.
    Definition T_rel_refl : forall x, T_rel x x := KI.T_rel_refl.
    Definition T_rel_trans : transitive T T_rel := KI.T_rel_trans.
  End Input.

  Module K := Knot_MixVariantHeredTOthRel(Input).
  Module KL := KnotLemmas_MixVariantHeredTOthRel(K).

  Definition knot: Type := K.knot.
  Definition ageable_knot : ageable knot := K.ageable_knot.
  #[global] Existing Instance ageable_knot.
  Definition knot_rel: knot -> knot -> Prop := K.knot_rel.
  Definition predicate: Type := KI.Pred knot ageable_knot knot_rel.

  Definition squash : (nat * KI.F predicate) -> knot :=
    fun k => K.squash (fst k, fmap KI.F KI.Pred2predicate KI.predicate2Pred (snd k)).

  Parameter unsquash : knot -> (nat * F predicate).

  Parameter approx : nat -> predicate -> predicate.

  Axiom squash_unsquash : forall k:knot, squash (unsquash k) = k.
  Axiom unsquash_squash : forall (n:nat) (f:F predicate),
    unsquash (squash (n,f)) = (n, fmap F (approx n) (approx n) f).

  Axiom approx_spec : forall n p ko,
    proj1_sig (Pred2predicate (approx n (predicate2Pred p))) ko =
     if (le_gt_dec n (level (fst ko))) then T_bot else proj1_sig p ko.

  Axiom knot_rel_spec: forall (k1 k2:knot),
     knot_rel k1 k2 =
    let (n,f) := unsquash k1 in
    let (n',f') := unsquash k2 in
    n = n' /\ Rel predicate f f'.

  Axiom knot_age1 : forall k:knot,
    age1 k =
    match unsquash k with
    | (O,_) => None
    | (S n,x) => Some (squash (n,x))
    end.

  Axiom knot_level : forall k:knot,
    level k = fst (unsquash k).

*)