(*
  Author(s):
    Andrej Dudenhefner (1)
  Affiliation(s):
    (1) Saarland University, Saarbrücken, Germany
*)

(*
  Reduction from:
    Finite Multiset Constraint Solvability (FMsetC_SAT)
  to:
    Linear Polynomial (over N) Constraint Solvability (LPolyNC_SAT)
*)

Require Import List PeanoNat Lia Permutation.
Import ListNotations.

Require Import Undecidability.SetConstraints.FMsetC.
Require Import Undecidability.PolynomialConstraints.LPolyNC.

Require Import Undecidability.Synthetic.Definitions.

Require Import ssreflect ssrbool ssrfun.

Set Default Goal Selector "!".

Module Argument.
Local Arguments poly_add !p !q.

Local Notation "p ≃ q" := (poly_eq p q) (at level 65).
Local Notation "A ≡ B" := (mset_eq A B) (at level 65).

Lemma poly_add_nthP {i p q} : nth i (poly_add p q) 0 = (nth i p 0) + (nth i q 0).
Proof. elim: i p q => [|i IH] [|a p [|b q]] /=; by [|lia]. Qed.

Definition encode_msetc (c : msetc) : polyc :=
  match c with
  | msetc_zero x => polyc_one x
  | msetc_sum x y z => polyc_sum x y z
  | msetc_h x y => polyc_prod x y
  end.

(* count the number of occurrences of each element *)
Fixpoint mset_to_poly (A: list nat) :=
  match A with
  | [] => []
  | a :: A => poly_add (repeat 0 a ++ [1]) (mset_to_poly A)
  end.

Lemma mset_to_poly_shift {a A B} : mset_to_poly (A ++ a :: B) ≃ mset_to_poly (a :: A ++ B).
Proof.
  elim: A; first done.
  move=> c A + i => /(_ i) /=. rewrite ?poly_add_nthP. by lia.
Qed.

Lemma mset_to_poly_appP {A B} : mset_to_poly (A ++ B) ≃ poly_add (mset_to_poly A) (mset_to_poly B).
Proof.
  elim: A; first done.
  move=> a A + i => /(_ i) /=. rewrite ?poly_add_nthP. by lia.
Qed.

Lemma mset_to_poly_mapP {A} : mset_to_poly (map S A) ≃ 0 :: mset_to_poly A.
Proof.
  elim: A; first by (case; [done | by case]).
  move=> a A + i => /(_ i). case: i.
  - by rewrite /map -/(map _ _) ?poly_add_nthP.
  - move=> i. rewrite /map -/(map _ _) /mset_to_poly -/(mset_to_poly).
    rewrite /= ?poly_add_nthP /=. by lia.
Qed.

Lemma poly_add_0I {p q r} : r ≃ [] -> p ≃ q -> p ≃ poly_add q r.
Proof.
  move=> + + i => /(_ i) + /(_ i). rewrite poly_add_nthP.
  case: i=> /=; by lia.
Qed.

Lemma poly_shiftI {p} : poly_mult [0; 1] p ≃ (0 :: p).
Proof.
  rewrite /poly_mult => i /=. rewrite poly_add_nthP /=.
  have ->: nth i (map (fun=> 0) p) 0 = 0. { by elim: i p => [|i IH] [|? p] /=. }
  move: i => [|i] /=; first done.
  rewrite poly_add_nthP.
  have ->: nth i [0] 0 = 0. { by move: i => [|[|i]]. }
  elim: i p => [|i IH] [|? p] /=; by [|lia].
Qed.

Lemma mset_to_poly_eqI {A B} : A ≡ B -> mset_to_poly A ≃ mset_to_poly B.
Proof.
  move=> /Permutation_count_occ. elim: A B.
  - by move=> B /Permutation_nil ->.
  - move=> a A IH B /[dup] /(Permutation_in a) /(_ (in_eq _ _)).
    move=> /(@in_split nat) [B1 [B2 ->]] /(@Permutation_cons_app_inv nat) /IH {}IH i.
    move: (IH i) (@mset_to_poly_shift a B1 B2 i).
    rewrite /= !poly_add_nthP. by lia.
Qed.

Lemma completeness {l} : FMsetC_SAT l -> LPolyNC_SAT (map encode_msetc l).
Proof.
  move=> [φ]. rewrite -Forall_forall => Hφ.
  exists (fun x => mset_to_poly (φ x)). rewrite -Forall_forall Forall_map.
  apply: Forall_impl; last by eassumption. case.
  - by move=> x /= /mset_to_poly_eqI.
  - move=> x y z /= /mset_to_poly_eqI + i => /(_ i) ->.
    by apply mset_to_poly_appP.
  - move=> x y /= /mset_to_poly_eqI + i => /(_ i) ->.
    move: (@poly_shiftI (mset_to_poly (φ y)) i) => /= ->.
    by apply: mset_to_poly_mapP.
Qed.

Fixpoint poly_to_mset (p: list nat) :=
  match p with
  | [] => []
  | a :: p => (repeat 0 a) ++ map S (poly_to_mset p)
  end.

Lemma count_occ_poly_to_msetP {a p}: count_occ Nat.eq_dec (poly_to_mset p) a = nth a p 0.
Proof.
  elim: a p.
  - case; first done.
    move=> + p /=. elim; first by elim: (poly_to_mset p).
    by move=> ? /= ->.
  - move=> i IH. case; first done.
    move=> a p /=. rewrite count_occ_app.
    rewrite -(count_occ_map S Nat.eq_dec) ?IH; first by lia.
    by elim a.
Qed.

Lemma poly_to_mset_eqI {p q} : p ≃ q -> poly_to_mset p ≡ poly_to_mset q.
Proof. move=> + a. by rewrite ?count_occ_poly_to_msetP. Qed.

Lemma poly_to_mset_addP {p q} : poly_to_mset (poly_add p q) ≡ poly_to_mset p ++ poly_to_mset q.
Proof. move=> a. by rewrite count_occ_app ? count_occ_poly_to_msetP poly_add_nthP. Qed.

Lemma poly_to_mset_consP {p} : poly_to_mset (0 :: p) = map S (poly_to_mset p).
Proof. done. Qed.

Lemma soundness {l} : LPolyNC_SAT (map encode_msetc l) -> FMsetC_SAT l.
Proof.
  have eq_trans A B C : A ≡ B -> B ≡ C -> A ≡ C.
  { move=> H1 H2 c. by rewrite (H1 c) (H2 c). }
  move=> [ψ]. rewrite -Forall_forall Forall_map => Hψ.
  exists (fun x => poly_to_mset (ψ x)). rewrite -Forall_forall.
  apply: Forall_impl; last by eassumption. case.
  - by move=> x /= /poly_to_mset_eqI.
  - move=> x y z /= /poly_to_mset_eqI. move /eq_trans. apply.
    by apply: poly_to_mset_addP.
  - move=> x y /= /poly_to_mset_eqI. move /eq_trans. apply.
    move: (ψ y) => p. rewrite -poly_to_mset_consP. apply: poly_to_mset_eqI.
    by apply: poly_shiftI.
Qed.

End Argument.

(* many-one reduction from FMsetC_SAT to LPolyNC_SAT *)
Theorem reduction : FMsetC_SAT ⪯ LPolyNC_SAT.
Proof.
  exists (map Argument.encode_msetc) => l. constructor.
  - exact Argument.completeness.
  - exact Argument.soundness.
Qed.


Require Coq.Logic.FunctionalExtensionality.
Require Import Coq.Program.Basics.
Require Import Coq.Arith.PeanoNat.
Require Import Omega.

Require Import Graph.

Open Scope program_scope.

(* An homomorphism is a morphism preserving the graph structure *)
Class Homomorphism {A B:Type} (f : Graph A -> Graph B) : Prop := {
  Hom_Empty :> f Empty = Empty;
  Hom_Overlay :> forall a b, f (Overlay a b) = Overlay (f a) (f b) ;
  Hom_Connect :> forall a b, f (Connect a b) = Connect (f a) (f b)
 }.

Lemma bind_is_hom {A B} (f: A -> Graph B): Homomorphism (bind f).
Proof.
  intros.
  repeat split.
Qed.

Lemma hom_is_bind {A B} {hom: Graph A -> Graph B}: Homomorphism hom -> hom = bind (hom ∘ Vertex).
Proof.
  intros H.
  apply FunctionalExtensionality.functional_extensionality.
  intro g.
  induction g.
  - compute. apply Hom_Empty.
  - auto.
  - rewrite Hom_Overlay.
    rewrite IHg1.
    rewrite IHg2.
    auto.
  - rewrite Hom_Connect.
    rewrite IHg1.
    rewrite IHg2.
    auto.
Qed.

(* Homomorphisms are exactly bind-functions *)
Theorem equiv_bind_hom {A B} (f : Graph A -> Graph B) :
  Homomorphism f <-> f = bind (f ∘ Vertex).
Proof.
  intros.
  split.
  - intro H.
    apply (hom_is_bind H).
  - intro H.
    rewrite H.
    apply bind_is_hom.
Qed.

(* The composition of a foldg function and a graph homorphism is a foldg-function *)
Theorem foldg_bind {A B C} {e : C} {v: B -> C} {o:C -> C -> C} {c:C -> C -> C} {f_v : A -> Graph B}:
  foldg e v o c ∘ bind f_v = foldg e (foldg e v o c ∘ f_v) o c.
Proof.
  apply FunctionalExtensionality.functional_extensionality.
  intro g.
  unfold compose.
  induction g.
  - auto.
  - auto.
  - pose (H :=bind_is_hom f_v).
    rewrite Hom_Overlay.
    repeat rewrite foldg_overlay.
    rewrite IHg1.
    rewrite IHg2.
    reflexivity.
  - pose (H :=bind_is_hom f_v).
    rewrite Hom_Connect.
    repeat rewrite foldg_connect.
    rewrite IHg1.
    rewrite IHg2.
    reflexivity.
Qed.

(* The composition of two graphs homomorphisms is a graph homomorphism *)
Theorem bind_compo {A B C} {hom1 : Graph A -> Graph B} {hom2 : Graph B -> Graph C}:
 (Homomorphism hom1) /\ (Homomorphism hom2) ->
  hom2 ∘ hom1 = bind (hom2 ∘ hom1 ∘ Vertex).
Proof.
  intros H.
  destruct H as (H1,H2).
  rewrite (hom_is_bind H2).
  rewrite (hom_is_bind H1).
  rewrite (inline_bind B C).
  rewrite foldg_bind.
  auto.
Qed.

Lemma size_ov {A} {g1 g2 : Graph A} : size (Overlay g1 g2) = size g1 + size g2.
Proof. auto. Qed.

Lemma size_co {A} {g1 g2 : Graph A} : size (Connect g1 g2) = size g1 + size g2.
Proof. auto. Qed.

Lemma sup1 (m n: nat) : (n >= 1) -> m + n >= 1.
Proof. omega. Qed.

Lemma size_sup1 A (g:Graph A) : size g >= 1.
Proof.
  induction g.
  - compute. auto.
  - compute. auto.
  - rewrite size_ov.
    apply sup1.
    exact IHg2.
  - rewrite size_co.
    apply sup1.
    exact IHg2.
Qed.

Lemma hom_leq_size A B (f: Graph A -> Graph B) : (Homomorphism f) -> forall g, size g <= size (f g).
Proof.
  intros H g.
  induction g.
  - rewrite Hom_Empty.
    auto.
  - apply size_sup1.
  - rewrite Hom_Overlay.
    repeat rewrite size_ov.
    omega.
  - rewrite Hom_Connect.
    repeat rewrite size_co.
    omega.
Qed.


From mathcomp Require Import all_ssreflect.

Implicit Type p q r: bool.
Implicit Type m n a b c: nat.

Lemma orbC p q: p || q = q || p.
  by case: p; case q. Qed.

Lemma Pirce p q: ((p ==> q) ==> p) ==> p.
  by case: p; case q. Qed.

Lemma bool_gimmics1 a: a != a.-1 -> a != 0.
  by elim: a. Qed.

Lemma find_me p q: ~~p = q -> p (+) q.
  by case: p; case q. Qed.

Lemma view_gimmics1 p a b: p -> (p ==> (a == b.*2)) -> a./2 = b.
  case p => //= _ => /eqP ->; exact: doubleK.
Qed.

Lemma view_gimmics2 p q r: ~~p && (r == q) -> q ==> (p || r).
  by case: p; case q; case r.
Qed.

Lemma iterSr A n (f: A -> A) x: iter n.+1 f x = iter n f (f x).
  elim: n => //= _ -> //.
Qed.

Lemma iter_predn m n: iter n predn m = m - n.
  elim: n => //=.
    by rewrite subn0.
    by move=> n ->; rewrite subnS.
Qed.

Lemma ltn_nqqAle m n: (m < n) = (m != n) && (m <= n).
    by rewrite ltnNge leq_eqVlt negb_or -leqNgt eq_sym.
Qed.

Lemma maxn_idPl m n: reflect (maxn m n = m) (m >= n).
  rewrite -subn_eq0 -(eqn_add2l m) addn0 maxnE.
  apply: eqP.
Qed.


Section InAHurry.

Lemma AndComm : forall a b : Prop, a /\ b -> b /\ a.
Proof.
intros a b H.
split.
  destruct H as [H1 H2].
  exact H2.
  destruct H as [H1 H2].
  exact H1.
Qed.

Lemma AndComm2 : forall a b : Prop, a /\ b -> b /\ a.
Proof.
intros a b H.
split; destruct H as [H1 H2].
  exact H2.
  exact H1.
Qed.

Lemma OrComm : forall a b : Prop, a \/ b -> b \/ a.
Proof.
intros a b H.
destruct H as [H1 | H2].
  right.
  exact H1.
  left.
  exact H2.
Qed.


(* infotheo: information theory and error-correcting codes in Coq             *)
(* Copyright (C) 2020 infotheo authors, license: LGPL-2.1-or-later            *)
From mathcomp Require Import all_ssreflect ssralg fingroup finalg matrix.
Require Import Reals.
From mathcomp Require Import Rstruct.
Require Import ssrR Reals_ext logb ssr_ext ssralg_ext bigop_ext Rbigop fdist.
Require Import proba entropy jfdist_cond.

(******************************************************************************)
(*                 Definition of channels and of the capacity                 *)
(*                                                                            *)
(*  `Ch(A, B) == discrete channel of input alphabet A and output alphabet B;  *)
(*               it is a collection of probability mass functions, one for    *)
(*               each a in A (i.e., a probability transition matrix           *)
(* `Ch*(A, B) == channels with non-empty alphabet                             *)
(* W `(b | a) == probability of receiving b knowing a was sent over the       *)
(*               channel W                                                    *)
(* W ``^ n, W ``(| x), W ``(y | x) == definition of a discrete memoryless     *)
(*               channel (DMC, or nth extension of the discrete memoryless    *)
(*               channel); W(y|x) = \Pi_i W_0(y_i|x_i) where W_0 is a         *)
(*               probability transition matrix                                *)
(*   `O(P, W) == output distribution for the channel                          *)
(* `H(P `o W) == output entropy for the channel                               *)
(* The input/output joint distribution for the channel is P `X W.             *)
(*  `H(P , W) == the input/output joint entropy for the channel               *)
(*  `H(W | P) == definition of conditional entropy using an input             *)
(*               distribution and a channel                                   *)
(*   `I(P, W) == the input/output mutual information for the channel          *)
(*   capacity == capacity of a channel                                        *)
(*                                                                            *)
(******************************************************************************)

Set Implicit Arguments.
Unset Strict Implicit.
Import Prenex Implicits.

Declare Scope channel_scope.
Delimit Scope fdist_scope with channel.
Local Open Scope channel_scope.

Reserved Notation "'`Ch(' A ',' B ')'" (at level 10, A, B at next level).
Reserved Notation "'`Ch*(' A ',' B ')'" (at level 10, A, B at next level).
Reserved Notation "W '`(' b '|' a ')'" (at level 10, b, a at next level).
Reserved Notation "W '``^' n" (at level 10).
Reserved Notation "W '``(|' x ')'" (at level 10, x at next level).
Reserved Notation "W '``(' y '|' x ')'" (at level 10, y, x at next level).
Reserved Notation "'`O(' P , W )" (at level 10, P, W at next level,
  format "'`O(' P ,  W )").
Reserved Notation "'`H(' P '`o' W )" (at level 10, P, W at next level,
  format "'`H(' P  '`o'  W )").
Reserved Notation "`H( P , W )" (at level 10, P, W at next level,
  format "`H( P ,  W )").
Reserved Notation "`H( W | P )" (at level 10, W, P at next level).
Reserved Notation "`I( P , W )" (at level 50, format "`I( P ,  W )").

Local Open Scope R_scope.

Module Channel1.
Section channel1.
Variables A B : finType.

Local Notation "'`Ch'" := (A -> fdist B) (only parsing).

Record chan_star := mkChan {
  c :> `Ch ;
  input_not_0 : (0 < #|A|)%nat }.

Local Notation "'`Ch*'" := (chan_star).

Lemma chan_star_eq (c1 c2 : `Ch*) : c c1 = c c2 -> c1 = c2.
Proof.
move: c1 c2 => [c1 Hc1] [c2 Hc2] /= <-{c2}; congr mkChan; exact: eq_irrelevance.
Qed.

End channel1.
End Channel1.

Definition chan_star_coercion := Channel1.c.
Coercion chan_star_coercion : Channel1.chan_star >-> Funclass.

Local Open Scope fdist_scope.
Local Open Scope proba_scope.

Notation "'`Ch(' A ',' B ')'" := (A -> {fdist B}) (only parsing) : channel_scope.
Notation "'`Ch*(' A ',' B ')'" := (@Channel1.chan_star A B) : channel_scope.
Notation "W '`(' b '|' a ')'" := ((W : `Ch(_, _)) a b) (only parsing) : channel_scope.

Local Open Scope proba_scope.
Local Open Scope vec_ext_scope.
Local Open Scope entropy_scope.

Module DMC.
Section def.
Local Open Scope ring_scope.
Variables (A B : finType) (W : `Ch(A, B)) (n : nat).

Definition f (x : 'rV[A]_n) :=
  [ffun y : 'rV[B]_n => (\prod_(i < n) W `(y ``_ i | x ``_ i))].

Lemma f0 x y : 0 <= f x y. Proof. rewrite ffunE; exact: prodR_ge0. Qed.

Lemma f1 x : (\sum_(y in 'rV_n) f x y = 1)%R.
Proof.
set f' := fun i b => W (x ``_ i) b.
suff H : (\sum_(g : {ffun 'I_n -> B}) \prod_(i < n) f' i (g i) = 1)%R.
  rewrite -{}[RHS]H /f'.
  rewrite (reindex_onto (fun vb : 'rV_n => [ffun x => vb ``_ x])
    (fun g  => \row_(k < n) g k)) /=; last first.
    move=> g _; apply/ffunP => /= i; by rewrite ffunE mxE.
  apply eq_big => vb.
  - rewrite inE.
    apply/esym/eqP/rowP => a; by rewrite mxE ffunE.
  - move=> _; rewrite ffunE; apply eq_bigr => i _; by rewrite ffunE.
by rewrite -bigA_distr_bigA /= /f' big1 // => i _; rewrite FDist.f1.
Qed.

Definition c : `Ch('rV[A]_n, 'rV[B]_n) :=
  locked (fun x => FDist.make (f0 x) (f1 x)).

End def.
End DMC.

Arguments DMC.c {A} {B}.

Notation "W '``^' n" := (DMC.c W n) : channel_scope.
Notation "W '``(|' x ')'" := (DMC.c W _ x) : channel_scope.
Notation "W '``(' y '|' x ')'" := (DMC.c W _ x y) : channel_scope.

Lemma DMCE (A B : finType) n (W : `Ch(A, B)) b a :
  W ``(b | a) = \prod_(i < n) W (a ``_ i) (b ``_ i).
Proof. by rewrite /DMC.c; unlock; rewrite ffunE. Qed.

Section DMC_sub_vec.
Variables (A B : finType) (W : `Ch(A, B)) (n : nat) (tb : 'rV[B]_n).

Lemma rprod_sub_vec (D : {set 'I_n}) (t : 'rV_n) :
  \prod_(i < #|D|) W ((t \# D) ``_ i) ((tb \# D) ``_ i) =
  \prod_(i in D) W (t ``_ i) (tb ``_ i).
Proof.
have [->|/set0Pn[i iD]] := eqVneq D set0.
  by rewrite big_set0 big_hasC //; apply/hasPn => /=; rewrite cards0; case.
pose f : 'I_n -> 'I_#|D| :=
  fun i => match Bool.bool_dec (i \in D) true with
             | left H => enum_rank_in H i
             | _ => enum_rank_in iD i
           end.
rewrite (reindex_onto (fun i : 'I_#|D| => enum_val i) f) /=.
  apply: eq_big => j; last by rewrite /sub_vec 2!mxE.
  rewrite /f /=; case: Bool.bool_dec => [a|].
    by rewrite enum_valK_in a eqxx.
  by rewrite enum_valP.
move=> j jD.
by rewrite /f /=; case: Bool.bool_dec => [a| //]; rewrite enum_rankK_in.
Qed.

Lemma DMC_sub_vecE (V : {set 'I_n}) (t : 'rV_n) :
  W ``(tb \# V | t \# V) = \prod_(i in V) W (t ``_ i) (tb ``_ i).
Proof. by rewrite DMCE -rprod_sub_vec. Qed.

End DMC_sub_vec.

Section fdist_out.
Variables (A B : finType) (P : fdist A) (W  : A -> fdist B).

Definition f := [ffun b : B => \sum_(a in A) W a b * P a].

Let f0 (b : B) : 0 <= f b.
Proof. by rewrite ffunE; apply: sumR_ge0 => a _; exact: mulR_ge0. Qed.

Let f1 : \sum_(b in B) f b = 1.
Proof.
under eq_bigr do rewrite ffunE /=.
rewrite exchange_big /= -(FDist.f1 P).
by apply eq_bigr => a _; rewrite -big_distrl /= (FDist.f1 (W a)) mul1R.
Qed.

Definition fdist_out : fdist B := locked (FDist.make f0 f1).

Lemma fdist_outE b : fdist_out b = \sum_(a in A) W a b * P a.
Proof. by rewrite /fdist_out; unlock; rewrite ffunE. Qed.

End fdist_out.

Notation "'`O(' P , W )" := (fdist_out P W) : channel_scope.

Notation "'`H(' P '`o' W )" := (`H ( `O( P , W ) )) : channel_scope.

Section fdist_out_prop.
Variables A B : finType.

Local Open Scope ring_scope.
Lemma fdist_rV_out (W : `Ch(A, B)) (P : fdist A) n (b : 'rV_n):
  `O(P, W) `^ _ b =
  \sum_(j : 'rV[A]_n) (\prod_(i < n) W j ``_ i b ``_ i) * P `^ _ j.
Proof.
rewrite fdist_rVE.
under eq_bigr do rewrite fdist_outE.
rewrite bigA_distr_big_dep /=.
rewrite (reindex_onto (fun p : 'rV_n => [ffun x => p ``_ x])
                      (fun y => \row_(k < n) y k)) //=; last first.
  by move=> i _; apply/ffunP => /= n0; rewrite ffunE mxE.
apply: eq_big.
- move=> a /=; apply/andP; split; first exact/finfun.familyP.
  by apply/eqP/rowP => a'; rewrite mxE ffunE.
- move=> a Ha; rewrite big_split /=; congr (_ * _)%R.
  + by apply eq_bigr => i /= _; rewrite ffunE.
  + by rewrite fdist_rVE; apply eq_bigr => i /= _; rewrite ffunE.
Qed.
Local Close Scope ring_scope.

Lemma fdistX_prod_out (W : `Ch(A, B)) (P : fdist A) : (fdistX (P `X W))`1 = `O(P, W).
Proof.
rewrite fdistX1; apply/fdist_ext => b; rewrite fdist_outE fdist_sndE.
by under eq_bigr do rewrite fdist_prodE mulRC.
Qed.

End fdist_out_prop.

Section Pr_fdist_prod.
Variables (A B : finType) (P : fdist A) (W : `Ch(A, B)) (n : nat).

Lemma Pr_DMC_rV_prod (Q : 'rV_n * 'rV_n -> bool) :
  Pr (((P `^ n) `X (W ``^ n))) [set x | Q x] =
  Pr ((P `X W) `^ n)           [set x | Q (rV_prod x)].
Proof.
rewrite /Pr [RHS]big_rV_prod /=; apply: eq_big => y.
  by rewrite !inE prod_rVK.
rewrite inE => Qy; rewrite fdist_prodE DMCE fdist_rVE -big_split /= fdist_rVE.
apply: eq_bigr => i /= _.
by rewrite fdist_prodE -snd_tnth_prod_rV -fst_tnth_prod_rV.
Qed.

Lemma Pr_DMC_fst (Q : 'rV_n -> bool) :
  Pr ((P `X W) `^ n) [set x | Q (rV_prod x).1 ] =
  Pr P `^ n          [set x | Q x].
Proof.
rewrite {1}/Pr big_rV_prod /= -(pair_big_fst _ _ [pred x | Q x]) //=; last first.
  move=> t /=.
  rewrite SetDef.pred_of_setE /= SetDef.finsetE /= ffunE. (* TODO: clean? *)
  congr (Q _).
  by apply/rowP => a; rewrite !mxE.
transitivity (\sum_(i | Q i) P `^ n i * (\sum_(y in 'rV[B]_n) W ``(y | i))).
  apply: eq_bigr => ta Sta; rewrite big_distrr; apply: eq_bigr => tb _ /=.
  rewrite DMCE [in RHS]fdist_rVE -[in RHS]big_split /= fdist_rVE.
  by apply eq_bigr => j _; rewrite fdist_prodE /= -fst_tnth_prod_rV -snd_tnth_prod_rV.
transitivity (\sum_(i | Q i) P `^ _ i).
  by apply eq_bigr => i _; rewrite (FDist.f1 (W ``(| i))) mulR1.
by rewrite /Pr; apply eq_bigl => t; rewrite !inE.
Qed.

Local Open Scope ring_scope.
Lemma Pr_DMC_out m (S : {set 'rV_m}) :
  Pr ((P `X W) `^ m) [set x | (rV_prod x).2 \notin S] =
  Pr (`O(P , W) `^ m) (~: S).
Proof.
rewrite {1}/Pr big_rV_prod /= -(pair_big_snd _ _ [pred x | x \notin S]) //=; last first.
  move=> tab /=.
  rewrite SetDef.pred_of_setE /= SetDef.finsetE /= ffunE. (* TODO: clean *)
  do 2 f_equal.
  by apply/rowP => a; rewrite !mxE.
rewrite /= /Pr /= exchange_big /=; apply: eq_big => tb; first by rewrite !inE.
move=> Htb.
rewrite fdist_rVE.
under [RHS]eq_bigr do rewrite fdist_outE.
rewrite bigA_distr_bigA /=.
rewrite (reindex_onto (fun p : 'rV[A]_m => [ffun x => p ord0 x])
    (fun y : {ffun 'I_m -> A} => \row_(i < m) y i)) /=; last first.
  by move=> f _; apply/ffunP => /= m0; rewrite ffunE mxE.
apply: eq_big => ta.
  by rewrite inE; apply/esym/eqP/rowP => a; rewrite mxE ffunE.
move=> Hta.
rewrite fdist_rVE /=; apply eq_bigr => l _.
by rewrite fdist_prodE -fst_tnth_prod_rV -snd_tnth_prod_rV ffunE mulRC.
Qed.
Local Close Scope ring_scope.

End Pr_fdist_prod.

Lemma channel_jcPr (A B : finType) (W : `Ch(A, B)) (P : fdist A) a b :
  P a != 0 ->
  W a b = \Pr_(fdistX (P `X W))[ [set b] | [set a] ].
Proof. by move=> Pa0; rewrite jcPr_fdistX_prod//; exact/eqP. Qed.

Notation "`H( P , W )" := (`H (P `X W)) : channel_scope.

Section conditional_entropy_chan.
Variables (A B : finType) (W : `Ch(A, B)) (P : fdist A).

Definition cond_entropy_chan := `H(P, W) - `H P.
End conditional_entropy_chan.

Notation "`H( W | P )" := (cond_entropy_chan W P) : channel_scope.

Section condentropychan_prop.
Variables (A B : finType) (W : `Ch(A, B)) (P : fdist A).

Lemma cond_entropy_chanE : `H(W | P) = cond_entropy (fdistX (P `X W)).
Proof.
rewrite /cond_entropy_chan.
have := chain_rule (P `X W); rewrite /joint_entropy => ->.
by rewrite fdist_prod1 addRC addRK.
Qed.

Lemma cond_entropy_chanE2 : `H(W | P) = \sum_(a in A) P a * `H (W a).
Proof.
rewrite cond_entropy_chanE cond_entropyE big_morph_oppR; apply: eq_bigr => a _.
rewrite big_morph_oppR /entropy mulRN -mulNR big_distrr/=; apply: eq_bigr => b _.
rewrite fdistXI fdist_prodE /= mulNR mulRA; congr (- _).
have [->|Pa0] := eqVneq (P a) 0; first by rewrite !(mulR0,mul0R).
by rewrite -channel_jcPr.
Qed.

End condentropychan_prop.

Section mutual_info_chan.
Local Open Scope fdist_scope.
Variables A B : finType.

Definition mutual_info_dist (P : {fdist A * B}) := `H P`1 + `H P`2 - `H P.

Definition mutual_info_chan P (W : `Ch(A, B)) := `H P + `H(P `o W) - `H(P , W).

End mutual_info_chan.

Notation "`I( P , W )" := (mutual_info_chan P W) : channel_scope.

Section mutual_info_chan_prop.
Variables (A B : finType) (W : `Ch(A, B)) (P : fdist A).

Lemma mutual_info_chanE : `I(P, W) = mutual_info (fdistX (P `X W)).
Proof.
rewrite /mutual_info_chan mutual_infoE -cond_entropy_chanE.
by rewrite -[in RHS]addR_opp oppRB addRCA addRA fdistX_prod_out.
Qed.

End mutual_info_chan_prop.

From mathcomp Require Import classical_sets.
Local Open Scope classical_set_scope.

Definition capacity (A B : finType) (W : `Ch(A, B)) :=
  reals.sup [set `I(P, W) | P in [set: fdist A]].


(** Simple Policy Optimizer based on Smart Constructors.
    This example nicely illustrates fully automatic axiomatic reasoning *)

Require Export NetKAT.

Module Optimize (F : FIELDSPEC) (V : VALUESPEC(F)).

Include NetKAT.NetKAT(F)(V).

Definition mk_union (p1 p2 : policy) :=
  match p1,p2 with
    | Drop, _ => p2
    | _, Drop => p1
    | _,_ => p1 + p2
  end.

Definition mk_seq (p1 p2 : policy) :=
  match p1,p2 with
    | Drop, _
    | _, Drop => Drop
    | Id, _ => p2
    | _, Id => p1
    | _, _ => p1;; p2
  end.

Fixpoint mk_star (p : policy) :=
  match p with
    | Drop
    | Id => Id
    | q* => mk_star q
    | _ => p*
  end.

Fixpoint optimize (p : policy) :=
  match p with
    | p1 + p2 => mk_union (optimize p1) (optimize p2)
    | p1;; p2 => mk_seq (optimize p1) (optimize p2)
    | p0* => mk_star (optimize p0)
    | _ => p
  end.

Lemma mk_union_sound: forall p q : policy, mk_union p q === p + q.
Proof. netkat_cases. Qed.

Lemma mk_seq_sound: forall p q : policy, mk_seq p q === p;;q.
Proof. netkat_cases. Qed.

Lemma star_zero: Drop* === Id.
Proof. rewrite <- ka_unroll_l. netkat. Qed.
Hint Rewrite star_zero : netkat.

Lemma star_one_aux: forall n h, eq (power n [|Id|] h) ([|Id|] h).
Proof.
  induction n; intros h h'; intuition.
  simpl in H. destruct H as [h'' [H0 H1]].
  + apply IHn in H1. congruence.
  + simpl. exists h. intuition. apply IHn. assumption.
Qed.

Lemma star_one: Id* === Id.
Proof.
  intros h h'.
  split; intros H.
  destruct H as [n].
  apply (star_one_aux n h). assumption.
  exists 0. simpl. apply H.
Qed.
Hint Rewrite star_one : netkat.

Lemma star_star p: p * * === p*.
Proof.
  intro h. split; intro H;
  destruct H as [n H]; generalize dependent x; generalize dependent h;
  induction n; intros.
  - exists 0. auto.
  - repeat destruct H.
    assert ([|p*|] x0 x) by eauto.
    destruct H1 as [m H1].
    exists ((x1 + m)%nat).
    apply power_decompose.
    exists x0; intuition.
  - exists 0. auto.
  - simpl in H. destruct H as [h'' [H0 H1]].
    assert (H2 := IHn _ _ H1); clear IHn H1 n.
    destruct H2 as [n]. exists (S n). simpl. exists h''.
    intuition. exists 1. simpl. exists h''; auto.
Qed.
Hint Rewrite star_star : netkat.

Lemma mk_star_sound p: mk_star p === p*.
Proof. netkat_induction p. Qed.

Hint Rewrite mk_union_sound mk_seq_sound mk_star_sound : netkat.

Theorem optimize_sound p: optimize p === p.
Proof. netkat_induction p. Qed.

End Optimize.


(************************************************************************)
(*  v      *   The Coq Proof Assistant  /  The Coq Development Team     *)
(* <O___,, *   INRIA - CNRS - LIX - LRI - PPS - Copyright 1999-2010     *)
(*   \VV/  **************************************************************)
(*    //   *      This file is distributed under the terms of the       *)
(*         *       GNU Lesser General Public License Version 2.1        *)
(************************************************************************)

(** Definition of a square root function for Z. *)

Require Import BinPos BinInt Psqrt.

Local Open Scope Z_scope.

Definition Zsqrtrem n :=
 match n with
  | 0 => (0, 0)
  | Zpos p =>
    match Psqrtrem p with
     | (s, IsPos r) => (Zpos s, Zpos r)
     | (s, _) => (Zpos s, 0)
    end
  | Zneg _ => (0,0)
 end.

Definition Zsqrt n :=
 match n with
  | 0 => 0
  | Zpos p => Zpos (Psqrt p)
  | Zneg _ => 0
 end.

Lemma Zsqrtrem_spec : forall n, 0<=n ->
 let (s,r) := Zsqrtrem n in n = s*s + r /\ 0 <= r <= 2*s.
Proof.
 destruct n. now repeat split.
 generalize (Psqrtrem_spec p). simpl.
 destruct 1; simpl; subst; now repeat split.
 now destruct 1.
Qed.

Lemma Zsqrt_spec : forall n, 0<=n ->
 let s := Zsqrt n in s*s <= n < (Zsucc s)*(Zsucc s).
Proof.
 destruct n. now repeat split. unfold Zsqrt.
 rewrite <- Zpos_succ_morphism. intros _. apply (Psqrt_spec p).
 now destruct 1.
Qed.

Lemma Zsqrt_neg : forall n, n<0 -> Zsqrt n = 0.
Proof.
 intros. now destruct n.
Qed.

Lemma Zsqrtrem_sqrt : forall n, fst (Zsqrtrem n) = Zsqrt n.
Proof.
 destruct n; try reflexivity.
 unfold Zsqrtrem, Zsqrt, Psqrt.
 destruct (Psqrtrem p) as (s,r). now destruct r.
Qed.


Require Import init.

Require Export mult_ring.

#[universes(template)]
Class Div U := {
    div : U → U
}.

Notation "/ a" := (div a) : algebra_scope.
Notation "a / b" := (a * /b) : algebra_scope.

Class MultLinv U `{Zero U, Mult U, One U, Div U} := {
    mult_linv : ∀ a, 0 ≠ a → /a * a = 1
}.
Class MultRinv U `{Zero U, Mult U, One U, Div U} := {
    mult_rinv : ∀ a, 0 ≠ a → a / a = 1
}.

Class FieldBase U `{
    FM : AllMult U,
    UD : Div U,
    UMD : @MultLinv U UZ UM UE UD,
    UMDR : @MultRinv U UZ UM UE UD
}.

Class Field U `{
    FF : FieldBase U,
    NotTrivial U
}.

Class HomomorphismDiv {U V} `{Zero U, Div U, Div V} (f : U → V) := {
    homo_div : ∀ a, 0 ≠ a → f (/a) = /f a
}.

(* begin hide *)
Arguments div : simpl never.

Section FieldImply1.

Context {U} `{Field U}.

Global Instance mult_linv_rinv : MultRinv U.
Proof.
    split.
    intros a a_nz.
    rewrite mult_comm.
    apply mult_linv.
    exact a_nz.
Qed.

End FieldImply1.

Section FieldImply2.

Context {U} `{Field U}.

Global Instance mult_linv_lcancel : MultLcancel U.
Proof.
    split.
    intros a b c c_nz eq.
    apply lmult with (/c) in eq.
    do 2 rewrite mult_assoc in eq.
    rewrite mult_linv in eq by exact c_nz.
    do 2 rewrite mult_lid in eq.
    exact eq.
Qed.

Global Instance mult_rinv_rcancel : MultRcancel U.
Proof.
    split.
    intros a b c c_nz eq.
    apply rmult with (/c) in eq.
    do 2 rewrite <- mult_assoc in eq.
    rewrite mult_rinv in eq by exact c_nz.
    do 2 rewrite mult_rid in eq.
    exact eq.
Qed.

End FieldImply2.

Section Field.

Context {U} `{Field U}.

(* end hide *)
Theorem div_nz : ∀ a, 0 ≠ a → 0 ≠ /a.
Proof.
    intros a a_nz eq.
    apply rmult with a in eq.
    rewrite mult_lanni in eq.
    rewrite mult_linv in eq by exact a_nz.
    exact (not_trivial_one eq).
Qed.

Theorem div_div : ∀ a, 0 ≠ a → /(/a) = a.
Proof.
    intros a a_nz.
    pose proof (div_nz a a_nz) as a'_nz.
    apply mult_lcancel with (/a); [>exact a'_nz|].
    rewrite mult_linv by exact a_nz.
    rewrite mult_rinv by exact a'_nz.
    reflexivity.
Qed.

Theorem mult_rrinv : ∀ a b, 0 ≠ b → a * b / b = a.
Proof.
    intros a b b_nz.
    rewrite <- mult_assoc.
    rewrite mult_rinv by exact b_nz.
    apply mult_rid.
Qed.
Theorem mult_rlinv : ∀ a b, 0 ≠ b → a / b * b = a.
Proof.
    intros a b b_nz.
    rewrite <- mult_assoc.
    rewrite mult_linv by exact b_nz.
    apply mult_rid.
Qed.
Theorem mult_lrinv : ∀ a b, 0 ≠ b → b * (/b * a) = a.
Proof.
    intros a b b_nz.
    rewrite mult_assoc.
    rewrite mult_rinv by exact b_nz.
    apply mult_lid.
Qed.
Theorem mult_llinv : ∀ a b, 0 ≠ b → /b * (b * a) = a.
Proof.
    intros a b b_nz.
    rewrite mult_assoc.
    rewrite mult_linv by exact b_nz.
    apply mult_lid.
Qed.

Theorem mult_llmove : ∀ a b c, 0 ≠ a → a * b = c ↔ b = /a * c.
Proof.
    intros a b c a_nz.
    split; intros eq.
    -   apply lmult with (/a) in eq.
        rewrite mult_llinv in eq by exact a_nz.
        exact eq.
    -   apply lmult with a in eq.
        rewrite mult_lrinv in eq by exact a_nz.
        exact eq.
Qed.
Theorem mult_lrmove : ∀ a b c, 0 ≠ b → a * b = c ↔ a = c / b.
Proof.
    intros a b c b_nz.
    split; intros eq.
    -   apply rmult with (/b) in eq.
        rewrite mult_rrinv in eq by exact b_nz.
        exact eq.
    -   apply rmult with b in eq.
        rewrite mult_rlinv in eq by exact b_nz.
        exact eq.
Qed.
Theorem mult_rlmove : ∀ a b c, 0 ≠ b → a = b * c ↔ /b * a = c.
Proof.
    intros a b c b_nz.
    split; intros eq.
    -   apply lmult with (/b) in eq.
        rewrite mult_llinv in eq by exact b_nz.
        exact eq.
    -   apply lmult with b in eq.
        rewrite mult_lrinv in eq by exact b_nz.
        exact eq.
Qed.
Theorem mult_rrmove : ∀ a b c, 0 ≠ c → a = b * c ↔ a / c = b.
Proof.
    intros a b c c_nz.
    split; intros eq.
    -   apply rmult with (/c) in eq.
        rewrite mult_rrinv in eq by exact c_nz.
        exact eq.
    -   apply rmult with c in eq.
        rewrite mult_rlinv in eq by exact c_nz.
        exact eq.
Qed.

Theorem mult_1_ab_da_b : ∀ a b, 0 ≠ a → 1 = a * b ↔ /a = b.
Proof.
    intros a b a_nz.
    rewrite mult_rlmove by exact a_nz.
    rewrite mult_rid.
    reflexivity.
Qed.
Theorem mult_1_ab_db_a : ∀ a b, 0 ≠ b → 1 = a * b ↔ /b = a.
Proof.
    intros a b b_nz.
    rewrite mult_rrmove by exact b_nz.
    rewrite mult_lid.
    reflexivity.
Qed.
Theorem mult_1_ab_a_db : ∀ a b, 0 ≠ b → 1 = a * b ↔ a = /b.
Proof.
    intros a b b_nz.
    rewrite mult_rrmove by exact b_nz.
    rewrite mult_lid.
    apply eq_iff.
Qed.
Theorem mult_1_ab_b_da : ∀ a b, 0 ≠ a → 1 = a * b ↔ b = /a.
Proof.
    intros a b a_nz.
    rewrite mult_rlmove by exact a_nz.
    rewrite mult_rid.
    apply eq_iff.
Qed.

Theorem mult_1_a_ab_b : ∀ a b, 0 ≠ b → 1 = a ↔ a * b = b.
Proof.
    intros a b b_nz.
    rewrite mult_lrmove by exact b_nz.
    rewrite mult_rinv by exact b_nz.
    apply eq_iff.
Qed.
Theorem mult_1_a_ba_b : ∀ a b, 0 ≠ b → 1 = a ↔ b * a = b.
Proof.
    intros a b b_nz.
    rewrite mult_llmove by exact b_nz.
    rewrite mult_linv by exact b_nz.
    apply eq_iff.
Qed.
Theorem mult_1_a_b_ab : ∀ a b, 0 ≠ b → 1 = a ↔ b = a * b.
Proof.
    intros a b b_nz.
    rewrite mult_rrmove by exact b_nz.
    rewrite mult_rinv by exact b_nz.
    reflexivity.
Qed.
Theorem mult_1_a_b_ba : ∀ a b, 0 ≠ b → 1 = a ↔ b = b * a.
Proof.
    intros a b b_nz.
    rewrite mult_rlmove by exact b_nz.
    rewrite mult_linv by exact b_nz.
    reflexivity.
Qed.

Theorem mult_1_dab_a_b : ∀ a b, 0 ≠ a → 1 = /a * b ↔ a = b.
Proof.
    intros a b a_nz.
    rewrite mult_1_ab_da_b by (apply div_nz; exact a_nz).
    rewrite div_div by exact a_nz.
    reflexivity.
Qed.
Theorem mult_1_adb_a_b : ∀ a b, 0 ≠ b → 1 = a / b ↔ a = b.
Proof.
    intros a b b_nz.
    rewrite mult_1_ab_a_db by (apply div_nz; exact b_nz).
    rewrite div_div by exact b_nz.
    reflexivity.
Qed.
Theorem mult_1_dab_b_a : ∀ a b, 0 ≠ a → 1 = /a * b ↔ b = a.
Proof.
    intros a b a_nz.
    rewrite mult_1_ab_b_da by (apply div_nz; exact a_nz).
    rewrite div_div by exact a_nz.
    reflexivity.
Qed.
Theorem mult_1_adb_b_a : ∀ a b, 0 ≠ b → 1 = a / b ↔ b = a.
Proof.
    intros a b b_nz.
    rewrite mult_1_ab_db_a by (apply div_nz; exact b_nz).
    rewrite div_div by exact b_nz.
    reflexivity.
Qed.

Theorem neg_div : ∀ a, 0 ≠ a → /(-a) = -/a.
Proof.
    intros a a_nz.
    pose proof (land (neg_nz _) a_nz) as na_nz.
    apply mult_rcancel with (-a); [>exact na_nz|].
    rewrite mult_linv by exact na_nz.
    rewrite mult_lneg, mult_rneg, neg_neg.
    rewrite mult_linv by exact a_nz.
    reflexivity.
Qed.

Theorem div_mult : ∀ a b, 0 ≠ a → 0 ≠ b → /(a * b) = /a * /b.
Proof.
    intros a b a_nz b_nz.
    apply mult_lcancel with a; [>exact a_nz|].
    rewrite mult_lrinv by exact a_nz.
    apply mult_lcancel with b; [>exact b_nz|].
    rewrite mult_rinv by exact b_nz.
    rewrite mult_assoc, (mult_comm b).
    rewrite mult_rinv; [>reflexivity|].
    apply mult_nz; assumption.
Qed.

Theorem div_one : /1 = 1.
Proof.
    rewrite <- (mult_lid (/1)).
    classic_case (0 = 1) as [triv|ntriv].
    -   rewrite <- triv.
        apply mult_lanni.
    -   rewrite mult_rinv by exact ntriv.
        reflexivity.
Qed.

(* begin hide *)
End Field.

Section MultHomo.

Context {U V} `{Field U, Field V}.
(* end hide *)
Context (f : U → V) `{
    @Injective U V f,
    @HomomorphismPlus U V UP UP0 f,
    @HomomorphismZero U V UZ UZ0 f,
    @HomomorphismNeg U V UN UN0 f,
    @HomomorphismMult U V UM UM0 f,
    @HomomorphismOne U V UE UE0 f,
    @HomomorphismDiv U V UZ UD UD0 f
}.

Global Instance field_inj : Injective f.
Proof.
    apply (homo_zero_inj _).
    intros a eq.
    classic_contradiction a_nz.
    apply (lmult (f (/a))) in eq.
    rewrite mult_ranni in eq.
    rewrite <- homo_mult in eq.
    rewrite mult_linv in eq by exact a_nz.
    rewrite homo_one in eq.
    contradiction (not_trivial_one eq).
Qed.
Local Remove Hints field_inj : typeclass_instances.

Global Instance field_homo_div : HomomorphismDiv f.
Proof.
    split.
    intros a a_nz.
    pose proof (inj_zero _ a_nz) as fa_nz.
    apply (mult_lcancel (f a) fa_nz).
    rewrite <- homo_mult.
    do 2 rewrite mult_rinv by assumption.
    apply homo_one.
Qed.
Local Remove Hints field_homo_div : typeclass_instances.

(* begin hide *)
End MultHomo.
(* end hide *)
Tactic Notation "mult_cancel_left" constr(x) :=
    mult_bring_left x;
    apply lmult.
Tactic Notation "mult_cancel_left" constr(x) "in" ident(H) :=
    mult_bring_left x in H;
    apply mult_lcancel in H.
Tactic Notation "mult_cancel_right" constr(x) :=
    mult_bring_right x;
    apply rmult.
Tactic Notation "mult_cancel_right" constr(x) "in" ident(H) :=
    mult_bring_right x in H;
    apply mult_rcancel in H.


Require Export SpecSyntax.
Set Implicit Arguments.

(******************************************************************************)
(* Shifting                                                                   *)
(******************************************************************************)

Fixpoint shiftIndex (c : nat) (i : nat) : nat :=
  match c with
    | O   => S i
    | S c =>
      match i with
        | O   => O
        | S i => S (shiftIndex c i)
      end
  end.

Fixpoint tshiftTy (c : nat) (T : Ty) : Ty :=
  match T with
    | tvar X      => tvar (shiftIndex c X)
    | top         => top
    | tarr T1 T2  => tarr (tshiftTy c T1) (tshiftTy c T2)
    | tall T1 T2  => tall (tshiftTy c T1) (tshiftTy (S c) T2)
    | tprod T1 T2 => tprod (tshiftTy c T1) (tshiftTy c T2)
  end.

Fixpoint tshiftTm (c : nat) (t : Tm) : Tm :=
  match t with
  | var x        => var x
  | abs T1 t2    => abs (tshiftTy c T1) (tshiftTm c t2)
  | app t1 t2    => app (tshiftTm c t1) (tshiftTm c t2)
  | tabs T1 t2   => tabs (tshiftTy c T1) (tshiftTm (S c) t2)
  | tapp t1 T2   => tapp (tshiftTm c t1) (tshiftTy c T2)
  | prod t1 t2   => prod (tshiftTm c t1) (tshiftTm c t2)
  | lett p t1 t2 => lett p (tshiftTm c t1) (tshiftTm c t2)
  end.

Fixpoint shiftTm (c : nat) (t : Tm) : Tm :=
  match t with
  | var x        => var (shiftIndex c x)
  | abs T1 t2    => abs T1 (shiftTm (S c) t2)
  | app t1 t2    => app (shiftTm c t1) (shiftTm c t2)
  | tabs T1 t2   => tabs T1 (shiftTm c t2)
  | tapp t1 T2   => tapp (shiftTm c t1) T2
  | prod t1 t2   => prod (shiftTm c t1) (shiftTm c t2)
  | lett p t1 t2 => lett p (shiftTm c t1) (shiftTm (bindPat p + c) t2)
  end.

Fixpoint weakenTm (t : Tm) (k : nat) : Tm :=
  match k with
    | O   => t
    | S k => shiftTm O (weakenTm t k)
  end.

Fixpoint tshiftExt (c : nat) (Δ : Ext) : Ext :=
  match Δ with
    | exempty    => exempty
    | exvar Δ T  => exvar (tshiftExt c Δ) (tshiftTy c T)
  end.

(******************************************************************************)
(* Type substitution.                                                         *)
(******************************************************************************)

Fixpoint tsubstIndex (X : nat) (T' : Ty) (Y : nat) : Ty :=
  match X , Y with
    | O   , O   => T'
    | O   , S Y => tvar Y
    | S X , O   => tvar O
    | S X , S Y => tshiftTy 0 (tsubstIndex X T' Y)
  end.

Fixpoint tsubstTy (X : nat) (T' : Ty) (T : Ty) : Ty :=
  match T with
    | tvar Y      => tsubstIndex X T' Y
    | top         => top
    | tarr T1 T2  => tarr (tsubstTy X T' T1) (tsubstTy X T' T2)
    | tall T1 T2  => tall (tsubstTy X T' T1) (tsubstTy (S X) T' T2)
    | tprod T1 T2 => tprod (tsubstTy X T' T1) (tsubstTy X T' T2)
  end.

Fixpoint tsubstTm (X : nat) (T' : Ty) (t : Tm) : Tm :=
  match t with
  | var x        => var x
  | abs T1 t2    => abs  (tsubstTy X T' T1) (tsubstTm X T' t2)
  | app t1 t2    => app  (tsubstTm X T' t1) (tsubstTm X T' t2)
  | tabs T1 t2   => tabs (tsubstTy X T' T1) (tsubstTm (S X) T' t2)
  | tapp t1 T2   => tapp (tsubstTm X T' t1) (tsubstTy X T' T2)
  | prod t1 t2   => prod (tsubstTm X T' t1) (tsubstTm X T' t2)
  | lett p t1 t2 => lett p (tsubstTm X T' t1) (tsubstTm X T' t2)
  end.

Fixpoint tsubstExt (X : nat) (T' : Ty) (Δ : Ext) : Ext :=
  match Δ with
    | exempty    => exempty
    | exvar Δ T  => exvar (tsubstExt X T' Δ) (tsubstTy X T' T)
  end.

(******************************************************************************)
(* Term substitutions.                                                        *)
(******************************************************************************)

Inductive Subst : Set :=
  | sub_here  : Subst
  | sub_var   : Subst → Subst
  | sub_bound : Subst → Subst.

Fixpoint weaken_subst (sub : Subst) (k : nat) : Subst :=
  match k with
   | O   => sub
   | S k => sub_var (weaken_subst sub k)
  end.

Fixpoint substIndex (x : Subst) (t : Tm) (y : nat) : Tm :=
  match x , y with
    | sub_here    , O   => t
    | sub_here    , S y => var y
    | sub_var x   , O   => var O
    | sub_var x   , S y => shiftTm O (substIndex x t y)
    | sub_bound x , y   => tshiftTm O (substIndex x t y)
  end.

Fixpoint substTm (x : Subst) (t' : Tm) (t : Tm) : Tm :=
  match t with
    | var y        => substIndex x t' y
    | abs T1 t2    => abs T1 (substTm (sub_var x) t' t2)
    | app t1 t2    => app (substTm x t' t1) (substTm x t' t2)
    | tabs T1 t2   => tabs T1 (substTm (sub_bound x) t' t2)
    | tapp t1 T2   => tapp (substTm x t' t1) T2
    | prod t1 t2   => prod (substTm x t' t1) (substTm x t' t2)
    | lett p t1 t2 => lett p (substTm x t' t1)
                        (substTm (weaken_subst x (bindPat p)) t' t2)
  end.

(******************************************************************************)
(* Context extension.                                                         *)
(******************************************************************************)

Fixpoint extend (Γ : Env) (Δ : Ext) : Env :=
  match Δ with
    | exempty     => Γ
    | exvar Δ T   => evar (extend Γ Δ) T
  end.

Fixpoint append (Δ1 Δ2 : Ext) : Ext :=
  match Δ2 with
    | exempty    => Δ1
    | exvar Δ2 T => exvar (append Δ1 Δ2) T
  end.

Fixpoint lengthExt (Δ : Ext) : nat :=
  match Δ with
    | exempty   => 0
    | exvar Δ _ => 1 + lengthExt Δ
  end.

(******************************************************************************)
(* Context lookups.                                                           *)
(******************************************************************************)

Inductive lookup_etvar : Env → nat → Ty → Prop :=
  | gb_here {Γ T} :
      lookup_etvar (etvar Γ T) O (tshiftTy O T)
  | gb_var {Γ T T' X} :
      lookup_etvar Γ X T →
      lookup_etvar (evar Γ T') X T
  | gb_bound {Γ T T' X} :
      lookup_etvar Γ X T →
      lookup_etvar (etvar Γ T') (S X) (tshiftTy O T).
Hint Constructors lookup_etvar.

Inductive lookup_evar : Env → nat → Ty → Prop :=
  | gv_here {Γ T} :
      lookup_evar (evar Γ T) O T
  | gv_var {Γ T T' X} :
      lookup_evar Γ X T →
      lookup_evar (evar Γ T') (S X) T
  | gv_bound {Γ T T' X} :
      lookup_evar Γ X T →
      lookup_evar (etvar Γ T') X (tshiftTy O T).
Hint Constructors lookup_evar.

(******************************************************************************)
(* Well-formedness.                                                           *)
(******************************************************************************)

(* These LoC shouldn't be counted here.. *)
Inductive wfTy (Γ: Env) : Ty → Prop :=
  | wf_tvar {X T} :
      lookup_etvar Γ X T → wfTy Γ (tvar X)
  | wf_top :
      wfTy Γ top
  | wf_tarr {T1 T2}:
      wfTy Γ T1 → wfTy Γ T2 → wfTy Γ (tarr T1 T2)
  | wf_tall {T1 T2} :
      wfTy Γ T1 → wfTy (etvar Γ T1) T2 →
      wfTy Γ (tall T1 T2)
  | wf_tprod {T1 T2}:
      wfTy Γ T1 → wfTy Γ T2 → wfTy Γ (tprod T1 T2).
Hint Constructors wfTy.


Require Import Coq.Lists.List Coq.Bool.Bool Coq.Structures.OrderedType Coq.Classes.Morphisms Coq.Setoids.Setoid.
Require Export Fiat.Common.Coq__8_4__8_5__Compat.

Definition SetEq {A: Type} (seq1: list A) (seq2: list A) :=
  forall x,
    List.In x seq1 <-> List.In x seq2.

Lemma SetEq_rewrite :
  forall {A: Type} (seq1 seq2: list A),
    SetEq seq1 seq2 <-> forall a, List.In a seq1 <-> List.In a seq2.
  unfold SetEq; tauto.
Qed.

Ltac autospecialize :=
  repeat match goal with
           | [ H: forall _, _ |- List.In ?x _ ] => try specialize (H x)
         end.

Lemma SetEq_Reflexive :
  forall {A: Type}, forall (x: list A), SetEq x x.
Proof.
  unfold Reflexive, SetEq;
  intuition;
  autospecialize;
  intuition.
Qed.

Lemma SetEq_Symmetric :
  forall {A: Type}, forall (x y: list A), SetEq x y -> SetEq y x.
Proof.
  unfold Symmetric, SetEq;
  intuition; autospecialize; intuition.
Qed.

Lemma SetEq_Transitive :
  forall {A: Type}, forall (x y z: list A), SetEq x y -> SetEq y z -> SetEq x z.
Proof.
  unfold Transitive, SetEq;
  intuition; autospecialize; intuition.
Qed.

Ltac seteq_equivalence :=
  eauto using SetEq_Transitive, SetEq_Symmetric, SetEq_Reflexive.

Lemma SetEq_Equivalence:
  forall {A: Type}, Equivalence (SetEq (A:=A)).
Proof.
  intros; constructor; seteq_equivalence.
Qed.

Lemma SetEq_Symmetric_iff:
  forall {A: Type}, forall (x y: list A), SetEq x y <-> SetEq y x.
Proof.
  split; seteq_equivalence.
Qed.

Lemma SetEq_trans_iff:
  forall {A: Type} (seq1 seq1' seq2: list A),
    SetEq seq1 seq1' ->
    (SetEq seq1 seq2 <-> SetEq seq1' seq2).
Proof.
  intuition; seteq_equivalence.
Qed.

Lemma SetEq_trans_iff_2:
  forall {A: Type} (seq1 seq2 seq2': list A),
    SetEq seq2 seq2' ->
    (SetEq seq1 seq2 <-> SetEq seq1 seq2').
Proof.
  intuition; seteq_equivalence.
Qed.

Definition SetUnion {A: Type} (x y: list A) := (x ++ y)%list.

Lemma union_left :
  forall {A: Type} (x: A) (seq1 seq2: list A),
    SetEq (SetUnion (x::seq1) seq2) (x :: (SetUnion seq1 seq2)).
Proof.
  intros; unfold SetEq, SetUnion; intuition.
Qed.

Lemma union_right :
  forall {A: Type} (x: A) (seq1 seq2: list A),
    SetEq (SetUnion seq1 (x::seq2)) (x :: (SetUnion seq1 seq2)).
Proof.
  intros; unfold SetEq, SetUnion; intuition;
  repeat (rewrite in_app_iff in *; simpl in *);
  intuition.
Qed.

Lemma filter_union :
  forall {A: Type} (seq1 seq2: list A),
  forall (pred: A -> bool),
    SetEq (List.filter pred (SetUnion seq1 seq2))
          (SetUnion (List.filter pred seq1) (List.filter pred seq2)).
Proof.
  unfold SetEq, SetUnion;
  split;
  intros;
  rewrite filter_In, in_app_iff in *;
  rewrite ! filter_In in *;
  tauto.
Qed.


Require Import Coq.ZArith.ZArith Coq.ZArith.BinIntDef Coq.ZArith.BinInt coqutil.Z.Lia.
Require Import coqutil.sanity coqutil.Tactics.forward coqutil.Word.Interface. Import word.
Require Import Kami.Lib.Word.
Require riscv.Utility.Utility.
From coqutil Require Import destr div_mod_to_equations.

Local Open Scope bool_scope.
Local Open Scope Z_scope.

Section KamiWordFacts.

  (* compatibility code for Coq 8.14 *)
  Fact Znat_N2Z_inj_mod : forall n m : N, m <> 0%N -> Z.of_N (n mod m) = Z.of_N n mod Z.of_N m.
  Proof. intros. now apply Znat.N2Z.inj_mod. Qed.

  Lemma wordToN_split2 a b w :
    wordToN (@split2 a b w) = BinNat.N.div (wordToN w) (NatLib.Npow2 a).
  Proof.
    pose proof wordToNat_split2 a b w as HH.
    eapply Nnat.Nat2N.inj_iff in HH.
    rewrite wordToN_nat, HH; f_equal; clear HH.
    rewrite wordToN_nat, NatLib.pow2_N.
    generalize (#w); intro.
    specialize (NatLib.zero_lt_pow2 a).
    generalize (NatLib.pow2 a); intros.
    pose proof Zdiv.div_Zdiv n n0 ltac:(blia).
    pose proof Znat.N2Z.inj_div (BinNat.N.of_nat n) (BinNat.N.of_nat n0).
    rewrite Znat.nat_N_Z in *.
    blia.
  Qed.

  Lemma wmsb_split2 a b w x y (H:b <> 0%nat)
    : wmsb (split2 a b w) x = wmsb w y.
  Proof.
    intros.
    rewrite <-(combine_split a b w) at 2.
    erewrite wmsb_combine by trivial.
    reflexivity.
  Qed.

  Lemma wordToZ_split2 a b w (H:b <> 0%nat)
    : wordToZ (@split2 a b w) = Z.div (wordToZ w) (2^Z.of_nat a).
  Proof.
    rewrite 2wordToZ_wordToN.
    rewrite wordToN_split2.
    erewrite wmsb_split2; [instantiate (1:=false)|trivial].
    case (wmsb w).
    all: rewrite ?Znat.N2Z.inj_div.
    all: rewrite <-?Znat.N_nat_Z, ?NatLib.Npow2_nat, ?N_Z_nat_conversions.Nat2Z.inj_pow.
    2: setoid_rewrite Z.add_0_r; trivial.
    rewrite ?Znat.Nat2Z.inj_add, ?Z.pow_add_r by blia.
    rewrite Z.mul_comm.
    symmetry.
    rewrite <-Z.add_opp_r, Zopp_mult_distr_l.
    rewrite Zdiv.Z_div_plus.
    1: reflexivity.
    eapply Z.lt_gt.
    eapply Z.pow_pos_nonneg; blia.
  Qed.

  Lemma wordToN_wones_ones:
    forall sz, wordToN (wones sz) = BinNat.N.ones (BinNat.N.of_nat sz).
  Proof.
    intros.
    rewrite wordToN_nat.
    rewrite wones_pow2_minus_one.
    rewrite Nnat.Nat2N.inj_sub, BinNat.N.sub_1_r.
    rewrite <-NatLib.pow2_N.
    cbv [BinNat.N.ones]; rewrite BinNat.N.shiftl_1_l.
    f_equal.
    apply Znat.N2Z.inj.
    rewrite NatLib.Z_of_N_Npow2, Znat.N2Z.inj_pow, Znat.nat_N_Z.
    reflexivity.
  Qed.

  Lemma uwordToZ_wplus_distr:
    forall sz (x y: Word.word sz),
      Z.of_N (wordToN (x ^+ y)) = (Z.of_N (wordToN x) + Z.of_N (wordToN y)) mod 2 ^ (Z.of_nat sz).
  Proof.
    intros.
    cbv [wplus wordBin].
    rewrite wordToN_NToWord_eqn, Znat_N2Z_inj_mod, Znat.N2Z.inj_add, NatLib.Z_of_N_Npow2.
    2: apply NatLib.Npow2_not_zero.
    f_equal; f_equal; blia.
  Qed.

  Lemma wnot_idempotent:
    forall {sz} (w: word sz),
      wnot (wnot w) = w.
  Proof.
    induction w; cbn; rewrite ?IHw, ?Bool.negb_involutive; eauto.
  Qed.

  Lemma wordToN_eq_rect:
    forall sz (w: Word.word sz) nsz Hsz,
      wordToN (eq_rect _ Word.word w nsz Hsz) = wordToN w.
  Proof.
    intros; subst; simpl; reflexivity.
  Qed.

  Lemma Z_pow_add_lor:
    forall n m p: Z,
      0 <= n < 2 ^ p -> 0 <= m -> 0 <= p ->
      (n + 2 ^ p * m)%Z = Z.lor n (2 ^ p * m).
  Proof.
    intros.
    apply eq_sym, BitOps.or_to_plus.
    rewrite Z.mul_comm, <-Z.shiftl_mul_pow2 by assumption.
    replace n with (Z.land n (Z.ones p)).
    - bitblast.Z.bitblast.
      rewrite Z.testbit_neg_r with (n:= l) by blia.
      apply Bool.andb_false_r.
    - destruct (Z.eq_dec n 0); [subst; apply Z.land_0_l|].
      assert (0 < n) by blia.
      rewrite Z.land_ones_low; [reflexivity|blia|].
      apply Z.log2_lt_pow2; blia.
  Qed.

  Lemma Z_of_wordToN_combine_alt:
    forall sz1 (w1: Word.word sz1) sz2 (w2: Word.word sz2),
      Z.of_N (wordToN (Word.combine w1 w2)) =
      Z.lor (Z.of_N (wordToN w1)) (Z.shiftl (Z.of_N (wordToN w2)) (Z.of_N (BinNat.N.of_nat sz1))).
  Proof.
    intros.
    rewrite wordToN_combine, N2Z.inj_add, N2Z.inj_mul.
    assert (0 <= Z.of_N (wordToN w1) < 2 ^ (Z.of_N (N.of_nat sz1))).
    { split; [apply N2Z.is_nonneg|].
      clear.
      induction w1; [simpl; blia|].
      unfold wordToN; fold wordToN.
      destruct b.
      { rewrite N2Z.inj_succ, N2Z.inj_mul, Nnat.Nat2N.inj_succ.
        rewrite N2Z.inj_succ.
        rewrite Z.pow_succ_r by blia; blia.
      }
      { rewrite N2Z.inj_mul, Nnat.Nat2N.inj_succ.
        rewrite N2Z.inj_succ.
        rewrite Z.pow_succ_r by blia; blia.
      }
    }
    assert (0 <= Z.of_N (wordToN w2)) by blia.
    assert (0 <= Z.of_N (N.of_nat sz1)) by blia.

    replace (Z.of_N (NatLib.Npow2 sz1)) with (Z.pow 2 (Z.of_N (N.of_nat sz1))).
    - generalize dependent (Z.of_N (wordToN w1)).
      generalize dependent (Z.of_N (wordToN w2)).
      generalize dependent (Z.of_N (N.of_nat sz1)).
      intros p ? z1 ? z2 ?.
      rewrite Z.shiftl_mul_pow2 by assumption.
      rewrite Z.mul_comm with (n:= z1).
      apply Z_pow_add_lor; assumption.
    - clear; induction sz1; [reflexivity|].
      rewrite Nnat.Nat2N.inj_succ, N2Z.inj_succ.
      unfold NatLib.Npow2; fold NatLib.Npow2.
      rewrite Z.pow_succ_r by blia; blia.
  Qed.

  Lemma ZToWord_zero:
    forall n, ZToWord n 0 = wzero n.
  Proof.
    destruct n; intros; [shatterer|].
    apply wordToZ_inj.
    rewrite wordToZ_ZToWord.
    - rewrite wordToZ_wzero; reflexivity.
    - split.
      + blia.
      + change 0 with (Z.of_nat 0).
        apply Nat2Z.inj_lt.
        apply NatLib.zero_lt_pow2.
  Qed.

  Lemma combine_wplus_wzero:
    forall sz1 (wb: Word.word sz1) sz2 (w1 w2: Word.word sz2),
      Word.combine wb w1 ^+ Word.combine (wzero sz1) w2 =
      Word.combine wb (w1 ^+ w2).
  Proof.
    induction wb; intros; [reflexivity|].
    simpl; rewrite <-wplus_WS_0.
    rewrite IHwb; reflexivity.
  Qed.

  Lemma split1_wplus_silent:
    forall sz1 sz2 (w1 w2: Word.word (sz1 + sz2)),
      split1 sz1 sz2 w2 = wzero _ ->
      split1 sz1 sz2 (w1 ^+ w2) = split1 sz1 sz2 w1.
  Proof.
    intros.
    pose proof (word_combinable _ _ w1).
    destruct H0 as [w11 [w12 ?]].
    pose proof (word_combinable _ _ w2).
    destruct H1 as [w21 [w22 ?]].
    subst; rewrite split1_combine in H; subst.
    rewrite combine_wplus_wzero.
    do 2 rewrite split1_combine.
    reflexivity.
  Qed.

  Lemma wordToN_split1 a b w :
    wordToN (@split1 a b w) = BinNat.N.modulo (wordToN w) (NatLib.Npow2 a).
  Proof.
    pose proof wordToNat_split1 a b w as HH.
    eapply Nnat.Nat2N.inj_iff in HH.
    rewrite wordToN_nat, HH; f_equal; clear HH.
    rewrite wordToN_nat, NatLib.pow2_N.
    generalize (#w); intro.
    remember (NatLib.pow2 a) as pa eqn:Ha.
    pose proof NatLib.pow2_zero a.
    pose proof mod_Zmod n pa ltac:(blia).
    pose proof Znat_N2Z_inj_mod (BinNat.N.of_nat n) (BinNat.N.of_nat pa) ltac:(blia).
    rewrite Znat.nat_N_Z in *.
    blia.
  Qed.

  Lemma sumbool_rect_weq {T} a b n x y :
    sumbool_rect (fun _ => T) (fun _ => a) (fun _ => b) (@weq n x y) = if weqb x y then a else b.
  Proof.
    cbv [sumbool_rect].
    destruct (weq _ _), (weqb _ _) eqn:?;
                                   try match goal with H : _ |- _ => eapply weqb_true_iff in H end;
      trivial; congruence.
  Qed.

  Lemma sumbool_rect_bool_weq n x y :
    sumbool_rect (fun _ => bool) (fun _ => true) (fun _ => false) (@weq n x y) = weqb x y.
  Proof. rewrite sumbool_rect_weq; destruct (weqb x y); trivial. Qed.

  Lemma unsigned_eqb n x y : Z.eqb (Z.of_N (wordToN x)) (Z.of_N (wordToN y)) = @weqb n x y.
  Proof.
    destruct (Z.eqb_spec (Z.of_N (wordToN x)) (Z.of_N (wordToN y))).
    - apply N2Z.inj, wordToN_inj in e; subst.
      apply eq_sym, weqb_eq; reflexivity.
    - apply eq_sym, weqb_ne.
      intro Hx; subst; auto.
  Qed.

  Lemma unsigned_split1_mod:
    forall n m w,
      Z.of_N (wordToN (split1 n m w)) = Z.of_N (wordToN w) mod (2 ^ (Z.of_nat n)).
  Proof.
    intros.
    rewrite wordToN_split1.
    rewrite N2Z.inj_mod by apply NatLib.Npow2_not_zero.
    rewrite NatLib.Z_of_N_Npow2.
    reflexivity.
  Qed.

End KamiWordFacts.

Section WithWidth.
  Context {width : Z}.
  Context {width_nonneg : Z.lt 0 width}.
  Local Notation sz := (Z.to_nat width).

  Definition kword: Type := Kami.Lib.Word.word sz.
  Definition kunsigned(x: kword): Z := Z.of_N (wordToN x).
  Definition ksigned: kword -> Z := @wordToZ sz.
  Definition kofZ: Z -> kword := ZToWord sz.

  Definition riscvZdivu(x y: Z): Z :=
    if y =? 0 then 2 ^ width - 1 else Z.div x y.

  Definition riscvZdivs(x y: Z): Z :=
    if (x =? - 2 ^ (width - 1)) && (y =? - 1) then x
    else if y =? 0 then - 1 else Z.quot x y.

  Definition riscvZmodu(x y: Z): Z :=
    if y =? 0 then x else Z.modulo x y.

  Definition riscvZmods(x y: Z): Z :=
    if y =? 0 then x else Z.rem x y.

  Instance word : word.word width := {|
    rep := kword;
    unsigned := kunsigned;
    signed := ksigned;
    of_Z := kofZ;

    add := @wplus sz;
    sub := @wminus sz;
    opp := @wneg sz;

    or  := @wor sz;
    and := @wand sz;
    xor := @wxor sz;
    not := @wnot sz;

    (* "x and not y" *)
    ndn x y := kofZ (Z.ldiff (kunsigned x) (kunsigned y));

    mul := @wmult sz;
    mulhss x y := kofZ (Z.mul (ksigned x) (ksigned y) / 2^width);
    mulhsu x y := kofZ (Z.mul (ksigned x) (kunsigned y) / 2^width);
    mulhuu x y := kofZ (Z.mul (kunsigned x) (kunsigned y) / 2^width);

    divu x y := kofZ (riscvZdivu (kunsigned x) (kunsigned y));
    divs x y := kofZ (riscvZdivs (ksigned x) (ksigned y));
    modu x y := kofZ (riscvZmodu (kunsigned x) (kunsigned y));
    mods x y := kofZ (riscvZmods (ksigned x) (ksigned y));

    (* shifts only look at the lowest 5-6 bits of the shift amount *)
    slu x y := wlshift x (Z.to_nat ((kunsigned y) mod width));
    sru x y := wrshift x (Z.to_nat ((kunsigned y) mod width));
    srs x y := wrshifta x (Z.to_nat ((kunsigned y) mod width));

    eqb := @weqb sz;
    ltu x y := if wlt_dec x y then true else false;
    lts x y := if wslt_dec x y then true else false;

    sextend oldwidth z := kofZ ((kunsigned z + 2^(oldwidth-1)) mod 2^oldwidth - 2^(oldwidth-1));

  |}.

  Section __.
    Import BinNat Word. Local Open Scope N_scope.

    Lemma wordToN_WS b n w :
      wordToN (@WS b n w) = 2*wordToN w + N.b2n b.
    Proof.
      case b; rewrite ?wordToN_WS_0, ?wordToN_WS_1; cbn [N.b2n].
      all : blia.
    Qed.

    Lemma testbit_wordToN_oob n (a : word n) i (H: Logic.not (i < N.of_nat n)) :
      N.testbit (wordToN a) i = false.
    Proof.
      pose proof wordToN_bound a.
      case (wordToN a) in *; trivial; intros.
      apply N.bits_above_log2, N.log2_lt_pow2; try blia.
      eapply N.lt_le_trans; try apply H0; clear H0.
      eapply Znat.N2Z.inj_le.
      rewrite NatLib.Z_of_N_Npow2, Znat.N2Z.inj_pow; cbn.
      eapply Z.pow_le_mono_r; blia.
    Qed.

    Lemma testbit_wordToN_bitwp_inbounds f n (a b : word n) i (H:i < N.of_nat n) :
      N.testbit (wordToN (bitwp f a b)) i = f (N.testbit (wordToN a) i) (N.testbit (wordToN b) i).
    Proof.
      revert dependent i; revert b; revert a; induction n; intros.
      { blia. }
      case (shatter_word_S a) as (?&?&?) in *; subst a.
      case (shatter_word_S b) as (?&?&?) in *; subst b.
      cbn [bitwp whd].
      rewrite 3wordToN_WS.
      case (N.eq_dec 0 i); intros.
      { subst. rewrite 3N.testbit_0_r; trivial. }
      { rewrite <-(N.succ_pred i) by blia.
        rewrite 3N.testbit_succ_r. eapply IHn. blia. }
    Qed.
  End __.

  Lemma uwordToZ_bitwp f F (F_spec : forall x y i, Z.testbit (F x y) i = f (Z.testbit x i) (Z.testbit y i)) sz (x y : Word.word sz)
    : uwordToZ (bitwp f x y) = (F (uwordToZ x) (uwordToZ y)) mod 2 ^ Z.of_nat sz.
  Proof.
    cbv [uwordToZ].
    eapply Z.bits_inj_iff'; intros.
    case (ZArith_dec.Z_lt_dec n (Z.of_nat sz)); intros.
    2: {
      rewrite Z.mod_pow2_bits_high by blia.
      rewrite ?Z.testbit_of_N' by trivial.
      rewrite testbit_wordToN_oob; trivial.
      intro X.
      eapply Znat.N2Z.inj_lt in X; blia.
    }
    rewrite Z.mod_pow2_bits_low by trivial.
    rewrite F_spec.
    rewrite ?Z.testbit_of_N' by trivial.
    rewrite testbit_wordToN_bitwp_inbounds; trivial.
    eapply Znat.N2Z.inj_lt; blia.
  Qed.

  Instance ok : word.ok word.
  Proof using width_nonneg.
    assert (AA: (0 < sz)%nat) by (eapply (Znat.Z2Nat.inj_lt 0); blia).
    assert (BB: Z.of_nat sz = width) by (rewrite Znat.Z2Nat.id; blia).
    split; trivial.
    all: cbv [rep unsigned signed of_Z add sub opp or and xor not
                  ndn mul mulhss mulhsu mulhuu divu divs modu mods slu sru srs
                  eqb ltu lts sextend word wrap
                  kword kunsigned ksigned kofZ]; intros.

    { pose proof @uwordToZ_ZToWord_full (Z.to_nat width) ltac:(blia).
      replace (Z.of_nat sz) with width in * by blia.
      match goal with H : _ |- _ => eapply H end. }
    { rewrite wordToZ_ZToWord_full; try blia. cbv [swrap].
      replace (Z.of_nat sz) with width in * by blia; trivial. }
    { rewrite ZToWord_Z_of_N, NToWord_wordToN; solve[trivial]. }
    { rewrite uwordToZ_wplus_distr, BB; reflexivity. }
    { cbv [wminus]; rewrite uwordToZ_wplus_distr, BB.
      destruct (BinNat.N.eq_dec (wordToN y) N0).
      { rewrite e.
        rewrite <-wordToN_wzero with (sz:= sz) in e.
        apply wordToN_inj in e; subst.
        rewrite wzero_wneg, wordToN_wzero; reflexivity.
      }
      { rewrite wneg_wordToN by assumption.
        rewrite Znat.N2Z.inj_sub by (pose proof (wordToN_bound y); blia).
        rewrite NatLib.Z_of_N_Npow2, BB.
        replace (Z.of_N (wordToN x) + (2 ^ width - Z.of_N (wordToN y)))
          with (Z.of_N (wordToN x) - Z.of_N (wordToN y) + 1 * 2 ^ width) by blia.
        rewrite Zdiv.Z_mod_plus_full.
        reflexivity.
      }
    }

    { destruct (BinNat.N.eq_dec (wordToN x) N0).
      { rewrite e.
        rewrite <-wordToN_wzero with (sz:= sz) in e.
        apply wordToN_inj in e; subst.
        rewrite wzero_wneg, wordToN_wzero; reflexivity.
      }
      { rewrite wneg_wordToN by assumption.
        rewrite Znat.N2Z.inj_sub by (pose proof (wordToN_bound x); blia).
        rewrite NatLib.Z_of_N_Npow2, BB.
        assert (Hms: Z.of_N (wordToN x) mod 2 ^ (Z.of_nat sz) = Z.of_N (wordToN x)).
        { apply Z.mod_small.
          split; [blia|].
          rewrite <-NatLib.Z_of_N_Npow2.
          apply Znat.N2Z.inj_lt, wordToN_bound.
        }
        rewrite BB in Hms.
        rewrite Zdiv.Z_mod_nz_opp_full by (rewrite Hms; blia).
        rewrite Hms; reflexivity.
      }
    }
    { setoid_rewrite (uwordToZ_bitwp _ _ Z.lor_spec); f_equal; congruence. }
    { setoid_rewrite (uwordToZ_bitwp _ _ Z.land_spec); f_equal; congruence. }
    { setoid_rewrite (uwordToZ_bitwp _ _ Z.lxor_spec); f_equal; congruence. }
    { rewrite wnot_wnot'_equiv. cbv [wnot'].
      setoid_rewrite (uwordToZ_bitwp _ _ Z.lxor_spec).
      rewrite <-Z.lxor_m1_l.
      pose proof uwordToZ_bound x.
      change (Z.of_N (wordToN x)) with (uwordToZ x).
      eapply Z.bits_inj_iff'; intros i Hi.
      case (ZArith_dec.Z_lt_dec i width); intros.
      2: rewrite !Z.mod_pow2_bits_high by blia; trivial.
      rewrite !Z.mod_pow2_bits_low by blia.
      rewrite 2Z.lxor_spec.
      rewrite bitblast.Z.testbit_minus1 by trivial.
      enough (Z.testbit (uwordToZ (wones sz)) i = true) by congruence.
      cbv [uwordToZ].
      rewrite ?Z.testbit_of_N' by trivial.
      rewrite wordToN_wones_ones.
      apply BinNat.N.ones_spec_low.
      blia.
    }

    { setoid_rewrite uwordToZ_ZToWord_full; f_equal; trivial; congruence. }
    { cbv [wmult wordBin].
      rewrite wordToN_NToWord_eqn, Znat_N2Z_inj_mod, Znat.N2Z.inj_mul, NatLib.Z_of_N_Npow2.
      2: apply NatLib.Npow2_not_zero.
      f_equal; f_equal; blia. }

    { rewrite wordToZ_ZToWord_full by blia;
        cbv [swrap];  repeat (blia || f_equal). }
    { rewrite wordToZ_ZToWord_full by blia;
        cbv [swrap];  repeat (blia || f_equal). }
    { repeat setoid_rewrite uwordToZ_ZToWord_full; try blia.
      cbv [swrap];  repeat (blia || f_equal). }

    { repeat setoid_rewrite uwordToZ_ZToWord_full; try blia.
      f_equal.
      2: repeat (blia || f_equal).
      (* f_equal. (* WHY (COQBUG?) does this add hyps to the goal *) *)
      cbv [riscvZdivu]; destr (Z.of_N (wordToN y) =? 0); blia. }
    { rewrite wordToZ_ZToWord_full by blia.
      cbv [swrap]; f_equal.
      2 : repeat (blia || f_equal).
      cbv [swrap]; f_equal.
      2 : repeat (blia || f_equal).
      cbv [riscvZdivs].
      destr ((wordToZ x =? - 2 ^ (width - 1))); cbn [andb].
      { destr (wordToZ y =? -1); try blia.
        destr (wordToZ y =?  0); try blia.
        f_equal; f_equal; blia. }
      { destr (wordToZ y =?  0); try blia.
        f_equal; f_equal; blia. } }
    { repeat setoid_rewrite uwordToZ_ZToWord_full; try blia.
      f_equal.
      2: repeat (blia || f_equal).
      cbv [riscvZmodu].
      destr (Z.of_N (wordToN y) =? 0); blia. }
    { rewrite wordToZ_ZToWord_full by blia.
      cbv [swrap]; f_equal.
      2 : repeat (blia || f_equal).
      cbv [swrap]; f_equal.
      2 : repeat (blia || f_equal).
      cbv [riscvZmods].
      destr (wordToZ y =? 0); try blia.
      f_equal; f_equal; blia. }
    { rewrite wlshift_mul_Zpow2 by (Z.div_mod_to_equations; blia).
      cbv [wmult wordBin].
      rewrite wordToN_NToWord_eqn, Znat_N2Z_inj_mod, Znat.N2Z.inj_mul, NatLib.Z_of_N_Npow2.
      2:apply NatLib.Npow2_not_zero.
      repeat setoid_rewrite uwordToZ_ZToWord_full; try blia.
      rewrite Zdiv.Zmult_mod_idemp_r.
      rewrite Z.shiftl_mul_pow2 by blia.
      f_equal. 2: { f_equal; blia. }
      f_equal.
      f_equal.
      apply Z.mod_small; blia. }

    { cbv [wrshift].
      rewrite wordToN_split2.
      cbv [eq_rec_r eq_rec].
      rewrite wordToN_nat, wordToNat_eq_rect, <-wordToN_nat, wordToN_combine, wordToN_wzero.
      PreOmega.zify.
      rewrite !NatLib.Z_of_N_Npow2.
      rewrite Z.shiftr_div_pow2 by blia.
      replace (Z.of_N (wordToN x) + 2 ^ Z.of_nat sz * 0)
        with (Z.of_N (wordToN x)) by blia.
      rewrite Z.mod_small by blia.
      rewrite Z.mod_small.
      2: {
        pose proof uwordToZ_bound x; cbv [uwordToZ] in *.
        replace (Z.of_nat sz) with width in * by blia.
        match goal with
        | H: _ \/ _ |- _ => clear H; subst
        end.
        pose proof Z.pow_pos_nonneg 2 (Z.of_N (wordToN y)) eq_refl. auto_specialize.
        replace 0 with (0/2 ^ Z.of_N (wordToN y)) by (apply Z.div_0_l; blia).
        split; eauto using Z.div_le_mono.
        eapply Z.div_lt_upper_bound; trivial.
        Lia.nia. }
      f_equal.
      f_equal.
      rewrite Znat.Z2Nat.id; blia. }

    { cbv [wrshifta eq_rec_r eq_rec].
      rewrite Z.mod_small, wordToZ_split2, wordToZ_eq_rect, sext_wordToZ, Znat.Z2Nat.id, Z.shiftr_div_pow2; try blia.
      cbv [swrap]; rewrite Z.mod_small; try blia.
      pose proof @wordToZ_size (pred sz).
      rewrite PeanoNat.Nat.succ_pred in H0; [|blia].
      specialize (H0 x).
      pose proof (wordToZ_size'' AA x); rewrite BB in H1.
      split.
      { assert (0 < 2 ^ Z.of_N (wordToN y)) by (apply Z.pow_pos_nonneg; blia).
        assert (0 < 2 ^ (width - 1)) by (apply Z.pow_pos_nonneg; blia).
        assert (- 2 ^ (width - 1) <= wordToZ x / 2 ^ Z.of_N (wordToN y)).
        { apply Z.div_le_lower_bound; [assumption|].
          etransitivity; [|apply H1].
          rewrite Z.mul_comm; apply Z.le_mul_diag_l; blia.
        }
        blia.
      }
      { apply Z.lt_add_lt_sub_r.
        replace (2 ^ width) with (2 * 2 ^ (width - 1)).
        2: { change 2 with (2 ^ 1) at 1.
             rewrite <-Z.pow_add_r by blia.
             f_equal; blia.
        }
        replace (2 * 2 ^ (width - 1) - 2 ^ (width - 1))
          with (2 ^ (width - 1)) by blia.
        apply Z.div_lt_upper_bound; [apply Z.pow_pos_nonneg; blia|].
        eapply Z.lt_le_trans; [apply H1|].
        rewrite Z.mul_comm; apply Z.le_mul_diag_r.
        { apply Z.pow_pos_nonneg; blia. }
        { pose proof (Z.pow_pos_nonneg 2 (Z.of_N (wordToN y))); blia. }
      }
    }
    { specialize (weqb_true_iff x y); case (weqb x y); intros [].
      { specialize (H eq_refl); subst; rewrite Z.eqb_refl; trivial. }
      { case (weq x y); try solve [intuition congruence]; intros HH.
        case (Z.eqb_spec (Z.of_N (wordToN x)) (Z.of_N (wordToN y))) as [X|X]; trivial.
        eapply Znat.N2Z.inj_iff in X.
        eapply wordToN_inj in X.
        contradiction. } }
    { case (wlt_dec x y) as [H|H]; cbv [wlt] in H;
        case (Z.ltb_spec (Z.of_N (wordToN x)) (Z.of_N (wordToN y)));
        trivial; blia. }
    { case (wslt_dec x y) as [H|H]; cbv [wslt] in H;
        case (Z.ltb_spec (wordToZ x) (wordToZ y)) as [G|G];
        trivial; blia. }
  Qed.
End WithWidth.
Arguments word : clear implicits.
Arguments ok : clear implicits.
Arguments kword : clear implicits.

#[global] Existing Instance word.
#[global] Existing Instance ok.

Open Scope Z_scope.

Section MkWords.
  Context {width : Z}.
  Context {width_cases : width = 32 \/ width = 64}.

  Lemma boundW: 0 < width.
  Proof.
    case width_cases; intro E; rewrite E; reflexivity.
  Defined.
  #[local] Instance wordW: word.word width := word width.
  #[local] Instance wordWok: word.ok wordW := ok width boundW.

  #[local] Instance word8: word.word 8 := word 8.
  #[local] Instance word8ok: word.ok word8 := ok 8 eq_refl.
End MkWords.


Require Import ssreflect.
From Tweetnacl.Libs Require Import Export.
From Tweetnacl.Gen Require Import Get_abcdef.

Section last_in.

Open Scope Z.

Context (U:Type).
Context (T:Type).
Context (T':Type).

Fixpoint fun_rec (f_last: U -> T) (f_in: nat -> T' -> U -> U) (m : nat) (z : T') (x : U) : T :=
  match m with
  | 0%nat => f_last x
  | S n => fun_rec f_last f_in n z (f_in n z x)
    end.

Fixpoint fun_rec_0 (f_in: nat -> T' -> U -> U) (m : nat) (z : T') (x : U) : U :=
  match m with
  | 0%nat => x
  | S n => fun_rec_0 f_in n z (f_in n z x)
   end.

Lemma fun_rec_extract_last: forall f_in f_last n z x,
  f_last (fun_rec_0 f_in n z x) = fun_rec f_last f_in  n z x.
Proof.
intros f_in f_last n.
induction n as [|n IHn]; intros z x.
reflexivity.
simpl.
rewrite IHn.
reflexivity.
Qed.

Close Scope Z.

End last_in.