(* Depth lower bounds (Section 5): no threshold unit, whatever its real weights, reads a memory word or
   computes a one-hot disjunction of conjunctions (Lemma 5.7), and no ternary unit computes the branch
   predicate on three or more bits, which one unit with integer weights computes (Remark 5.5). *)
From Stdlib Require Import Reals Lra ZArith Lia List Bool.
Import ListNotations.

(* ---- units with real weights ------------------------------------------------------------------------------ *)
(* the pre-activation without bias of the unit with weights w on the inputs x_i, i in l *)
Fixpoint rsum (w : nat -> R) (x : nat -> bool) (l : list nat) : R :=
  match l with [] => 0%R | i :: t => ((if x i then w i else 0) + rsum w x t)%R end.

Definition rfire (N : nat) (w : nat -> R) (b : R) (x : nat -> bool) : bool :=
  if Rle_dec 0 (rsum w x (seq 0 N) + b) then true else false.

(* two inputs whose ones, counted with multiplicity, are those of two others have the same summed sums *)
Lemma rsum_swap : forall w x1 x2 x3 x4 l,
  (forall i, (Nat.b2n (x1 i) + Nat.b2n (x2 i) = Nat.b2n (x3 i) + Nat.b2n (x4 i))%nat) ->
  (rsum w x1 l + rsum w x2 l = rsum w x3 l + rsum w x4 l)%R.
Proof.
  intros w x1 x2 x3 x4 l H. induction l as [|i t IH]; cbn [rsum]; [lra|].
  specialize (H i). destruct (x1 i), (x2 i), (x3 i), (x4 i); cbn [Nat.b2n] in H; try lia; lra.
Qed.

Lemma rfire_swap : forall N w b x1 x2 x3 x4,
  (forall i, (Nat.b2n (x1 i) + Nat.b2n (x2 i) = Nat.b2n (x3 i) + Nat.b2n (x4 i))%nat) ->
  rfire N w b x1 = true -> rfire N w b x2 = true -> rfire N w b x3 = false -> rfire N w b x4 = false -> False.
Proof.
  intros N w b x1 x2 x3 x4 H. unfold rfire. pose proof (rsum_swap w x1 x2 x3 x4 (seq 0 N) H).
  destruct (Rle_dec 0 (rsum w x1 (seq 0 N) + b)), (Rle_dec 0 (rsum w x2 (seq 0 N) + b)),
    (Rle_dec 0 (rsum w x3 (seq 0 N) + b)), (Rle_dec 0 (rsum w x4 (seq 0 N) + b));
    intros; try discriminate; lra.
Qed.

(* ---- Lemma 5.7 --------------------------------------------------------------------------------------------- *)
(* the address alpha in the inputs 0 .. A - 1, least significant first, the words mu in the inputs A .. *)
Fixpoint bval (x : nat -> bool) (k : nat) : nat :=
  match k with O => 0 | S k' => bval x k' + Nat.b2n (x k') * 2 ^ k' end.

Definition select (A : nat) (x : nat -> bool) : bool := x (A + bval x A)%nat.

Lemma bval_low : forall x k, (forall i, (1 <= i)%nat -> x i = false) -> bval x k = match k with O => O | S _ => Nat.b2n (x 0%nat) end.
Proof.
  intros x k H. induction k as [|k IH]; [reflexivity|]. cbn [bval]. rewrite IH.
  destruct k as [|k]; [cbn; lia|]. rewrite (H (S k)) by lia. cbn [Nat.b2n]. lia.
Qed.

Theorem read_needs_two_levels : forall A w b, (1 <= A)%nat ->
  ~ (forall x, rfire (A + 2 ^ A) w b x = select A x).
Proof.
  intros A w b HA H.
  set (pt := fun (a m0 m1 : bool) i => (a && Nat.eqb i 0) || (m0 && Nat.eqb i A) || (m1 && Nat.eqb i (S A))).
  assert (Hlow : forall a m0 m1 i, (1 <= i)%nat -> (i < A)%nat -> pt a m0 m1 i = false).
  { intros a m0 m1 i H1 H2. unfold pt. destruct (Nat.eqb_spec i 0), (Nat.eqb_spec i A), (Nat.eqb_spec i (S A));
      try lia; rewrite !andb_false_r; reflexivity. }
  assert (Hsel : forall a m0 m1, select A (pt a m0 m1) = if a then m1 else m0).
  { intros a m0 m1. unfold select.
    assert (E : bval (pt a m0 m1) A = Nat.b2n a).
    { assert (G : forall k, (k <= A)%nat -> bval (pt a m0 m1) k = match k with O => O | S _ => Nat.b2n a end).
      { induction k as [|k IH]; intros Hk; [reflexivity|]. cbn [bval]. rewrite IH by lia.
        destruct k as [|k].
        - unfold pt. rewrite Nat.eqb_refl. destruct (Nat.eqb_spec 0 A), (Nat.eqb_spec 0 (S A)); try lia.
          rewrite !andb_false_r, !orb_false_r, andb_true_r. destruct a; reflexivity.
        - rewrite (Hlow a m0 m1 (S k)) by lia. cbn [Nat.b2n]. lia. }
      rewrite (G A (le_n _)). destruct A; [lia | reflexivity]. }
    rewrite E. unfold pt. destruct a; cbn [Nat.b2n].
    - replace (A + 1)%nat with (S A) by lia. rewrite Nat.eqb_refl.
      destruct (Nat.eqb_spec (S A) 0), (Nat.eqb_spec (S A) A); try lia. rewrite !andb_false_r, andb_true_r. reflexivity.
    - rewrite Nat.add_0_r, Nat.eqb_refl. destruct (Nat.eqb_spec A 0), (Nat.eqb_spec A (S A)); try lia.
      rewrite !andb_false_r, andb_true_r, orb_false_r. reflexivity. }
  apply (rfire_swap (A + 2 ^ A) w b (pt false true false) (pt true false true) (pt false false true) (pt true true false)).
  - intro i. unfold pt. destruct (Nat.eqb_spec i 0), (Nat.eqb_spec i A), (Nat.eqb_spec i (S A));
      cbn [andb orb Nat.b2n]; lia.
  - rewrite H, Hsel. reflexivity.
  - rewrite H, Hsel. reflexivity.
  - rewrite H, Hsel. reflexivity.
  - rewrite H, Hsel. reflexivity.
Qed.

(* the one-hot form: the lines s in the inputs 0 .. n - 1, the bits mu in the inputs n .. 2n - 1 *)
Definition onehot (n : nat) (x : nat -> bool) : Prop := exists j, (j < n)%nat /\ forall i, (i < n)%nat -> x i = Nat.eqb i j.

Definition orsel (n : nat) (x : nat -> bool) : bool := existsb (fun j => x j && x (n + j)%nat) (seq 0 n).

Theorem onehot_needs_two_levels : forall n w b, (2 <= n)%nat ->
  ~ (forall x, onehot n x -> rfire (n + n) w b x = orsel n x).
Proof.
  intros n w b Hn H.
  set (pt := fun (s m : nat) i => Nat.eqb i s || Nat.eqb i (n + m)%nat).
  assert (Hoh : forall s m, (s < n)%nat -> onehot n (pt s m)).
  { intros s m Hs. exists s. split; [exact Hs|]. intros i Hi. unfold pt.
    destruct (Nat.eqb_spec i (n + m)); [lia|]. apply orb_false_r. }
  assert (Hval : forall s m, (s < 2)%nat -> (m < 2)%nat -> orsel n (pt s m) = Nat.eqb s m).
  { intros s m Hs Hm. unfold orsel, pt.
    destruct (Nat.eqb_spec s m) as [<-|E].
    - apply existsb_exists. exists s. split; [apply in_seq; lia|]. rewrite !Nat.eqb_refl, orb_true_l, orb_true_r.
      reflexivity.
    - apply not_true_is_false. intro Hx. apply existsb_exists in Hx as (j & Hj & Hx). apply in_seq in Hj.
      apply andb_true_iff in Hx as [H1 H2]. apply orb_true_iff in H1, H2.
      destruct H1 as [H1|H1], H2 as [H2|H2]; apply Nat.eqb_eq in H1, H2; lia. }
  apply (rfire_swap (n + n) w b (pt 0 0) (pt 1 1) (pt 0 1) (pt 1 0))%nat.
  - intro i. unfold pt. destruct (Nat.eqb_spec i 0), (Nat.eqb_spec i 1), (Nat.eqb_spec i (n + 0)),
      (Nat.eqb_spec i (n + 1)); cbn; lia.
  - rewrite H by (apply Hoh; lia). rewrite Hval by lia. reflexivity.
  - rewrite H by (apply Hoh; lia). rewrite Hval by lia. reflexivity.
  - rewrite H by (apply Hoh; lia). rewrite Hval by lia. reflexivity.
  - rewrite H by (apply Hoh; lia). rewrite Hval by lia. reflexivity.
Qed.

(* ---- Remark 5.5 ------------------------------------------------------------------------------------------- *)
Open Scope Z_scope.

Fixpoint zsum (w : nat -> Z) (x : nat -> bool) (l : list nat) : Z :=
  match l with [] => 0 | i :: t => (if x i then w i else 0) + zsum w x t end.

Definition zfire (n : nat) (w : nat -> Z) (b : Z) (x : nat -> bool) : bool := Z.leb 0 (zsum w x (seq 0 n) + b).

(* the value of r_0 .. r_(n-1) in two's complement, and the branch predicate [value <= 0] *)
Fixpoint uval (r : nat -> bool) (k : nat) : Z :=
  match k with O => 0 | S k' => uval r k' + Z.b2z (r k') * 2 ^ Z.of_nat k' end.

Definition tval (n : nat) (r : nat -> bool) : Z := uval r (n - 1) - Z.b2z (r (n - 1)%nat) * 2 ^ Z.of_nat (n - 1).

Definition branch (n : nat) (r : nat -> bool) : bool := Z.leb (tval n r) 0.

(* with integer weights one unit computes it *)
Lemma branch_one_unit : forall n r,
  branch n r = Z.leb 0 (Z.b2z (r (n - 1)%nat) * 2 ^ Z.of_nat (n - 1) - uval r (n - 1)).
Proof. intros n r. unfold branch, tval. apply eq_true_iff_eq. rewrite !Z.leb_le. lia. Qed.

Lemma zsum_ext : forall w x x' l, (forall i, In i l -> x i = x' i) -> zsum w x l = zsum w x' l.
Proof.
  intros w x x' l H. induction l as [|i t IH]; cbn [zsum]; [reflexivity|].
  rewrite (H i (or_introl eq_refl)), IH by (intros j Hj; apply H; right; exact Hj). reflexivity.
Qed.

Lemma zsum_zero : forall w l, zsum w (fun _ => false) l = 0.
Proof. intros w l. induction l as [|i t IH]; cbn [zsum]; lia. Qed.

Lemma zsum_one : forall w k l, NoDup l -> In k l -> zsum w (fun i => Nat.eqb i k) l = w k.
Proof.
  intros w k l Hnd. induction Hnd as [|i t Hi Hnd IH]; intros Hk; [destruct Hk|]. cbn [zsum].
  destruct (Nat.eqb_spec i k) as [->|E].
  - rewrite (zsum_ext w _ (fun _ => false)), zsum_zero; [lia|].
    intros j Hj. destruct (Nat.eqb_spec j k); [subst; contradiction | reflexivity].
  - destruct Hk as [Hk|Hk]; [contradiction|]. rewrite IH by exact Hk. lia.
Qed.

Lemma zsum_all : forall w l, zsum w (fun _ => true) l = fold_right Z.add 0 (map w l).
Proof. intros w l. induction l as [|i t IH]; cbn [zsum map fold_right]; lia. Qed.

Lemma uval_ext : forall r r' k, (forall i, (i < k)%nat -> r i = r' i) -> uval r k = uval r' k.
Proof.
  intros r r' k H. induction k as [|k IH]; cbn [uval]; [reflexivity|].
  rewrite IH by (intros i Hi; apply H; lia). rewrite (H k) by lia. reflexivity.
Qed.

Lemma uval_one : forall j k, (j < k)%nat -> uval (fun i => Nat.eqb i j) k = 2 ^ Z.of_nat j.
Proof.
  intros j k Hjk. induction k as [|k IH]; [lia|]. cbn [uval].
  destruct (Nat.eqb_spec k j) as [->|E].
  - rewrite (uval_ext _ (fun _ => false)).
    + assert (Z0 : forall m, uval (fun _ => false) m = 0) by (induction m as [|m IHm]; cbn [uval]; [reflexivity | rewrite IHm; reflexivity]).
      rewrite Z0. cbn [Z.b2z]. lia.
    + intros i Hi. destruct (Nat.eqb_spec i j); [lia | reflexivity].
  - rewrite IH by lia. cbn [Z.b2z]. lia.
Qed.

Lemma uval_all : forall k, uval (fun _ => true) k = 2 ^ Z.of_nat k - 1.
Proof.
  induction k as [|k IH]; cbn [uval]; [reflexivity|]. rewrite IH. cbn [Z.b2z].
  rewrite Nat2Z.inj_succ, Z.pow_succ_r by lia. lia.
Qed.

Lemma sum_ternary_lower : forall (w : nat -> Z) (l : list nat), (forall i, In i l -> w i = -1) -> fold_right Z.add 0 (map w l) = - Z.of_nat (length l).
Proof.
  intros w l H. induction l as [|i t IH]; cbn [map fold_right length]; [reflexivity|].
  rewrite (H i (or_introl eq_refl)), IH by (intros j Hj; apply H; right; exact Hj). lia.
Qed.

Theorem branch_not_ternary : forall n w b, (3 <= n)%nat -> (forall i, -1 <= w i <= 1) ->
  ~ (forall r, zfire n w b r = branch n r).
Proof.
  intros n w b Hn Hw H.
  (* r = 0 has value 0, so b >= 0 *)
  assert (Hb : 0 <= b).
  { specialize (H (fun _ => false)). unfold zfire, branch, tval in H. rewrite zsum_zero in H.
    assert (Z0 : forall m, uval (fun _ => false) m = 0) by (induction m as [|m IHm]; cbn [uval]; [reflexivity | rewrite IHm; reflexivity]).
    rewrite Z0 in H. cbn [Z.b2z] in H. replace (0 - 0 * 2 ^ Z.of_nat (n - 1)) with 0 in H by lia.
    cbn in H. apply Z.leb_le. rewrite H. reflexivity. }
  (* each unit vector e_i, i < n - 1, has value 2^i > 0, so w_i = -1 and b = 0 *)
  assert (He : forall i, (i < n - 1)%nat -> w i + b < 0).
  { intros i Hi. specialize (H (fun k => Nat.eqb k i)). unfold zfire, branch, tval in H.
    rewrite zsum_one in H by (apply seq_NoDup || (apply in_seq; lia)).
    rewrite uval_one in H by exact Hi. destruct (Nat.eqb_spec (n - 1) i) as [E|E]; [lia|]. cbn [Z.b2z] in H.
    assert (P : 0 < 2 ^ Z.of_nat i) by (apply Z.pow_pos_nonneg; lia).
    replace (2 ^ Z.of_nat i - 0 * 2 ^ Z.of_nat (n - 1) <=? 0) with false in H
      by (symmetry; apply Z.leb_gt; lia).
    apply Z.leb_gt. exact H. }
  assert (Hw1 : forall i, (i < n - 1)%nat -> w i = -1) by (intros i Hi; specialize (He i Hi); specialize (Hw i); lia).
  assert (Hb0 : b = 0) by (specialize (He 0%nat ltac:(lia)); rewrite Hw1 in He by lia; lia).
  (* the vector of ones has value -1, so w_(n-1) >= n - 1 >= 2 *)
  specialize (H (fun _ => true)). unfold zfire, branch, tval in H.
  rewrite zsum_all, uval_all in H. cbn [Z.b2z] in H.
  replace (seq 0 n) with (seq 0 (n - 1) ++ [(n - 1)%nat]) in H
    by (replace n with (S (n - 1)) at 3 by lia; rewrite seq_S; reflexivity).
  rewrite map_app, fold_right_app in H. cbn [map fold_right] in H.
  assert (Hs : forall (l : list Z) c, fold_right Z.add c l = fold_right Z.add 0 l + c).
  { intros l c. induction l as [|y t IH]; cbn [fold_right]; lia. }
  rewrite Hs, sum_ternary_lower, length_seq in H by (intros i Hi; apply in_seq in Hi; apply Hw1; lia).
  replace (2 ^ Z.of_nat (n - 1) - 1 - 1 * 2 ^ Z.of_nat (n - 1) <=? 0) with true in H by (symmetry; apply Z.leb_le; lia).
  apply Z.leb_le in H. specialize (Hw (n - 1)%nat). lia.
Qed.

(* ---- Lemma 5.4 --------------------------------------------------------------------------------------------- *)
(* z(r) = H(-(r_0 + ... + r_(n-1))) *)
Definition zunit (n : nat) (r : nat -> bool) : bool := zfire n (fun _ => -1) 0 r.

Lemma zsum_neg_le : forall r l, zsum (fun _ => -1) r l <= 0.
Proof. intros r l. induction l as [|i t IH]; cbn [zsum]; [lia|]. destruct (r i); lia. Qed.

Lemma zsum_neg : forall r l, Z.leb 0 (zsum (fun _ => -1) r l) = forallb (fun i => negb (r i)) l.
Proof.
  intros r l. induction l as [|i t IH]; cbn [zsum forallb]; [reflexivity|].
  pose proof (zsum_neg_le r t). destruct (r i); cbn [negb andb].
  - apply Z.leb_gt. lia.
  - rewrite Z.add_0_l. exact IH.
Qed.

Lemma uval_bound : forall r k, 0 <= uval r k < 2 ^ Z.of_nat k.
Proof.
  intros r k. induction k as [|k IH]; cbn [uval]; [cbn; lia|].
  rewrite Nat2Z.inj_succ, Z.pow_succ_r by lia. destruct (r k); cbn [Z.b2z]; lia.
Qed.

Lemma uval_zero : forall r k, Z.eqb (uval r k) 0 = forallb (fun i => negb (r i)) (seq 0 k).
Proof.
  intros r k. induction k as [|k IH]; [reflexivity|]. cbn [uval]. rewrite seq_S, forallb_app, <- IH.
  cbn [forallb Nat.add]. pose proof (uval_bound r k) as B.
  assert (P : 0 < 2 ^ Z.of_nat k) by (apply Z.pow_pos_nonneg; lia).
  destruct (r k); cbn [Z.b2z negb andb].
  - rewrite andb_false_r. apply Z.eqb_neq. lia.
  - rewrite andb_true_r, Z.mul_0_l, Z.add_0_r. reflexivity.
Qed.

(* z(r) is 1 exactly when r = 0, and [r <= 0] = H(r_(n-1) + z(r) - 1): the branch takes two ternary units *)
Theorem branch_unit : forall n r, (1 <= n)%nat ->
  zunit n r = Z.eqb (tval n r) 0 /\ branch n r = Z.leb 0 (Z.b2z (r (n - 1)%nat) + Z.b2z (zunit n r) - 1).
Proof.
  intros n r Hn.
  assert (Hz : zunit n r = Z.eqb (uval r (n - 1)) 0 && negb (r (n - 1)%nat)).
  { unfold zunit, zfire. rewrite Z.add_0_r, zsum_neg.
    replace (seq 0 n) with (seq 0 (n - 1) ++ [(n - 1)%nat])
      by (replace n with (S (n - 1)) at 3 by lia; rewrite seq_S; reflexivity).
    rewrite forallb_app, <- uval_zero. cbn [forallb]. rewrite andb_true_r. reflexivity. }
  unfold branch, tval. rewrite Hz. pose proof (uval_bound r (n - 1)) as B.
  assert (P : 0 < 2 ^ Z.of_nat (n - 1)) by (apply Z.pow_pos_nonneg; lia).
  set (u := uval r (n - 1)) in *.
  destruct (r (n - 1)%nat); cbn [negb]; rewrite ?andb_false_r, ?andb_true_r; cbn [Z.b2z].
  - split; [symmetry; apply Z.eqb_neq; lia|].
    rewrite (proj2 (Z.leb_le (u - 1 * 2 ^ Z.of_nat (n - 1)) 0)) by lia. reflexivity.
  - replace (u - 0 * 2 ^ Z.of_nat (n - 1)) with u by lia. split; [reflexivity|].
    destruct (Z.eqb_spec u 0) as [E|E]; cbn [Z.b2z]; [rewrite E; reflexivity|].
    rewrite (proj2 (Z.leb_gt u 0)) by lia. reflexivity.
Qed.