(* Levelization (Lemma 2.5): a clocked netlist whose units have maximal depth d >= 1 is computed by a layered map
   of depth d.  Each unit that contributes to a next value occurs in it with its weights and bias, once if its
   depth is less than d and once for each state bit whose next value it is if its depth is d, and every other row
   is an identity unit. *)
From Stdlib Require Import List Bool Arith ZArith Lia.
Import ListNotations.
From TH Require Import TwoHead Net.

(* ---- layered maps (Definition 2.3) ----------------------------------------------------------------------- *)
(* a layer: rows reading the outputs of the layer before by position *)
Definition lfire (L : list tu) (y : list bool) : list bool := map (fun u => fire u (fun i => nth i y false)) L.
Definition lmap (Ls : list (list tu)) (x : list bool) : list bool := fold_left (fun y L => lfire L y) Ls x.

(* the rows of each layer read positions of the layer before, distinct ones, with weights -1 or 1 *)
Definition rowok (w : nat) (u : tu) : Prop :=
  NoDup (map fst (tin u)) /\ forall i z, In (i, z) (tin u) -> i < w /\ (z = 1 \/ z = -1)%Z.
Fixpoint lwf (w : nat) (Ls : list (list tu)) : Prop :=
  match Ls with [] => True | L :: t => Forall (rowok w) L /\ lwf (length L) t end.

Lemma lmap_app : forall Ls L x, lmap (Ls ++ [L]) x = lfire L (lmap Ls x).
Proof. intros Ls L x. unfold lmap. rewrite fold_left_app. reflexivity. Qed.

(* ---- depth -------------------------------------------------------------------------------------------- *)
Fixpoint maxl (l : list nat) : nat := match l with [] => 0 | x :: t => Nat.max x (maxl t) end.

Lemma maxl_ge : forall l a, In a l -> a <= maxl l.
Proof. induction l as [|x t IH]; intros a H; [destruct H|]. destruct H as [<-|H]; cbn [maxl]; [lia|]. specialize (IH a H). lia. Qed.

Lemma maxl_le : forall l c, (forall a, In a l -> a <= c) -> maxl l <= c.
Proof.
  induction l as [|x t IH]; intros c H; cbn [maxl]; [lia|].
  pose proof (H x (or_introl eq_refl)). pose proof (IH c (fun a Ha => H a (or_intror Ha))). lia.
Qed.

Lemma maxl_pos : forall l, 1 <= maxl l -> exists a, In a l /\ 1 <= a.
Proof.
  induction l as [|x t IH]; cbn [maxl]; intro H; [lia|].
  destruct (Nat.le_gt_cases 1 x) as [Hx|Hx]; [exists x; split; [left; reflexivity | exact Hx]|].
  destruct (IH ltac:(lia)) as (a & Ha & H1). exists a. split; [right; exact Ha | exact H1].
Qed.

(* the depth of a signal: 0 for a state bit, one more than its deepest predecessor for a unit *)
Fixpoint depf (us : list tu) (m : nat) (D : nat -> nat) : nat -> nat :=
  match us with
  | [] => D
  | u :: t => depf t (S m) (fun i => if Nat.eqb i m then S (maxl (map (fun p => D (fst p)) (tin u))) else D i)
  end.

Lemma depf_prefix : forall us m D i, i < m -> depf us m D i = D i.
Proof.
  induction us as [|u t IH]; intros m D i H; cbn [depf]; [reflexivity|].
  rewrite IH by lia. destruct (Nat.eqb_spec i m); [lia | reflexivity].
Qed.

Lemma depf_at : forall us m D j, j < length us ->
  (forall k, k < length us -> reads_below (nth k us tu0) (m + k)) ->
  depf us m D (m + j) = S (maxl (map (fun p => depf us m D (fst p)) (tin (nth j us tu0)))).
Proof.
  induction us as [|u t IH]; intros m D j Hj Hwf; cbn [length] in Hj, Hwf; [lia|].
  destruct j as [|j].
  - cbn [nth depf]. rewrite Nat.add_0_r, depf_prefix by lia. rewrite Nat.eqb_refl. f_equal. f_equal.
    apply map_ext_in. intros [i z] Hin. cbn [fst].
    pose proof (Hwf 0 ltac:(lia) i z Hin) as Hi. cbn [nth] in Hi.
    rewrite depf_prefix by lia. destruct (Nat.eqb_spec i m); [lia | reflexivity].
  - cbn [nth depf]. replace (m + S j) with (S m + j) by lia. apply IH; [lia|].
    intros k Hk i z Hin. specialize (Hwf (S k) ltac:(lia) i z Hin). lia.
Qed.

Lemma map_nth_self : forall (l : list nat) d, map (fun i => nth i l d) (seq 0 (length l)) = l.
Proof.
  induction l as [|a t IH]; intro d; [reflexivity|]. cbn [length seq map nth]. f_equal.
  rewrite <- seq_shift, map_map. apply IH.
Qed.

Section Lev.
Variable n : nat.
Variable us : list tu.
Variable nx : list nat.

Definition NS : nat := n + length us.
Definition wfn : Prop := forall k, k < length us -> reads_below (nth k us tu0) (n + k).

Definition dep : nat -> nat := depf us n (fun _ => 0).
Definition dmax : nat := maxl (map (fun k => dep (n + k)) (seq 0 (length us))).

(* unit k reads signal s *)
Definition reads (k s : nat) : bool := existsb (fun p => Nat.eqb (fst p) s) (tin (nth k us tu0)).

Lemma reads_in : forall k s, reads k s = true <-> exists z, In (s, z) (tin (nth k us tu0)).
Proof.
  intros k s. unfold reads. rewrite existsb_exists. split.
  - intros ([i z] & Hin & E). cbn [fst] in E. apply Nat.eqb_eq in E. subst i. exists z. exact Hin.
  - intros (z & Hin). exists (s, z). split; [exact Hin | apply Nat.eqb_refl].
Qed.

Lemma dep_state : forall i, i < n -> dep i = 0.
Proof. intros i H. unfold dep. rewrite depf_prefix by exact H. reflexivity. Qed.

Lemma dep_unit : wfn -> forall k, k < length us ->
  dep (n + k) = S (maxl (map (fun p => dep (fst p)) (tin (nth k us tu0)))).
Proof. intros Hwf k Hk. unfold dep. apply depf_at; assumption. Qed.

Lemma dep_pred : wfn -> forall k p z, k < length us -> In (p, z) (tin (nth k us tu0)) -> dep p < dep (n + k).
Proof.
  intros Hwf k p z Hk Hin. rewrite dep_unit by assumption. apply Nat.lt_succ_r.
  apply maxl_ge. apply in_map_iff. exists (p, z). split; [reflexivity | exact Hin].
Qed.

Lemma dep_le_dmax : forall k, k < length us -> dep (n + k) <= dmax.
Proof. intros k Hk. apply maxl_ge. apply in_map_iff. exists k. split; [reflexivity | apply in_seq; lia]. Qed.

Lemma dep_unit_pos : wfn -> forall k, k < length us -> 1 <= dep (n + k).
Proof. intros Hwf k Hk. rewrite dep_unit by assumption. lia. Qed.

Lemma dep_zero : wfn -> forall s, s < NS -> dep s = 0 -> s < n.
Proof.
  intros Hwf s Hs H. destruct (Nat.lt_ge_cases s n) as [Hl|Hg]; [exact Hl|].
  pose proof (dep_unit_pos Hwf (s - n) ltac:(unfold NS in Hs; lia)). replace (n + (s - n)) with s in H0 by lia. lia.
Qed.

Lemma dep_le : forall s, s < NS -> dep s <= dmax.
Proof.
  intros s Hs. destruct (Nat.lt_ge_cases s n) as [Hl|Hg]; [rewrite dep_state by exact Hl; lia|].
  replace s with (n + (s - n)) by lia. apply dep_le_dmax. unfold NS in Hs. lia.
Qed.

(* ---- the signals some next value depends on -------------------------------------------------------- *)
Fixpoint needf (vs : list tu) (m : nat) (S0 : nat -> bool) : nat -> bool :=
  match vs with
  | [] => S0
  | u :: t => let N' := needf t (S m) S0 in
              fun i => N' i || (N' m && existsb (fun p => Nat.eqb (fst p) i) (tin u))
  end.

Definition need : nat -> bool := needf us n (fun i => existsb (Nat.eqb i) nx).

Lemma needf_seed : forall vs m S0 i, S0 i = true -> needf vs m S0 i = true.
Proof. induction vs as [|u t IH]; intros m S0 i H; cbn [needf]; [exact H|]. rewrite IH by exact H. reflexivity. Qed.

Lemma needf_closed : forall vs m S0 k p z, k < length vs ->
  (forall j, j < length vs -> reads_below (nth j vs tu0) (m + j)) ->
  needf vs m S0 (m + k) = true -> In (p, z) (tin (nth k vs tu0)) -> needf vs m S0 p = true.
Proof.
  induction vs as [|u t IH]; intros m S0 k p z Hk Hwf Hn Hin; cbn [length] in Hk; [lia|].
  assert (Hu : forall i z', In (i, z') (tin u) -> i < m) by (intros i z' H; specialize (Hwf 0 ltac:(cbn; lia) i z' H); cbn [nth] in Hwf; lia).
  assert (Hm : existsb (fun p0 => Nat.eqb (fst p0) (m + k)) (tin u) = true -> False).
  { intro E. apply existsb_exists in E as ([i z'] & Hi & E). cbn [fst] in E. apply Nat.eqb_eq in E. subst i.
    specialize (Hu _ _ Hi). lia. }
  cbn [needf] in *. destruct k as [|k].
  - cbn [nth] in Hin. rewrite Nat.add_0_r in Hn. apply orb_true_iff in Hn as [Hn|Hn].
    + apply orb_true_iff. right. rewrite Hn. cbn [andb]. apply existsb_exists. exists (p, z). split; [exact Hin|].
      apply Nat.eqb_refl.
    + apply andb_true_iff in Hn as [Hn _]. apply orb_true_iff. right. rewrite Hn. cbn [andb].
      apply existsb_exists. exists (p, z). split; [exact Hin | apply Nat.eqb_refl].
  - cbn [nth] in Hin. apply orb_true_iff in Hn as [Hn|Hn].
    + apply orb_true_iff. left. replace (m + S k) with (S m + k) in Hn by lia.
      apply (IH (S m) S0 k p z); [lia | | exact Hn | exact Hin].
      intros j Hj. specialize (Hwf (S j) ltac:(cbn; lia)). cbn [nth] in Hwf. replace (S m + j) with (m + S j) by lia. exact Hwf.
    + apply andb_true_iff in Hn as [_ Hn]. exfalso. apply Hm. replace (m + k) with (m + k) by lia.
      replace (m + S k) with (m + S k) in Hn by lia. apply existsb_exists in Hn as ([i z'] & Hi & E).
      cbn [fst] in E. apply Nat.eqb_eq in E. subst i. specialize (Hu _ _ Hi). lia.
Qed.

Lemma needf_step : forall vs m S0 i, needf vs m S0 i = true ->
  S0 i = true \/ exists k z, k < length vs /\ needf vs m S0 (m + k) = true /\ In (i, z) (tin (nth k vs tu0)).
Proof.
  induction vs as [|u t IH]; intros m S0 i H; cbn [needf] in *; [left; exact H|].
  apply orb_true_iff in H as [H|H].
  - destruct (IH (S m) S0 i H) as [H0|(k & z & Hk & Hn & Hin)]; [left; exact H0|].
    right. exists (S k), z. cbn [length nth]. split; [lia|]. split; [|exact Hin].
    replace (m + S k) with (S m + k) by lia. rewrite Hn. reflexivity.
  - apply andb_true_iff in H as [H1 H2]. apply existsb_exists in H2 as ([i' z] & Hi & E). cbn [fst] in E.
    apply Nat.eqb_eq in E. subst i'. right. exists 0, z. cbn [length nth]. split; [lia|]. split; [|exact Hi].
    rewrite Nat.add_0_r, H1. reflexivity.
Qed.

(* contributing to a next value: being one, or a predecessor of a unit that contributes *)
Inductive contrib : nat -> Prop :=
| contrib_out : forall s, In s nx -> contrib s
| contrib_pred : forall k s z, k < length us -> contrib (n + k) -> In (s, z) (tin (nth k us tu0)) -> contrib s.

Lemma need_nx : forall s, In s nx -> need s = true.
Proof.
  intros s H. apply needf_seed. apply existsb_exists. exists s. split; [exact H | apply Nat.eqb_refl].
Qed.

Lemma need_pred : wfn -> forall k p z, k < length us -> need (n + k) = true -> In (p, z) (tin (nth k us tu0)) ->
  need p = true.
Proof. intros Hwf k p z Hk Hn Hin. exact (needf_closed us n _ k p z Hk Hwf Hn Hin). Qed.

Theorem need_contrib : wfn -> forall s, need s = true <-> contrib s.
Proof.
  intros Hwf s. split.
  - (* by induction on the distance to the last signal *)
    assert (G : forall d s, NS - s <= d -> need s = true -> contrib s \/ NS <= s).
    { induction d as [|d IH]; intros s' Hd Hn.
      - right. lia.
      - destruct (needf_step us n _ s' Hn) as [H0|(k & z & Hk & Hk2 & Hin)].
        + left. apply contrib_out. apply existsb_exists in H0 as (y & Hy & E). apply Nat.eqb_eq in E. subst. exact Hy.
        + assert (Hs : s' < n + k) by (exact (Hwf k Hk s' z Hin)).
          destruct (IH (n + k) ltac:(unfold NS in *; lia) Hk2) as [Hc|Hc]; [|unfold NS in Hc; lia].
          left. exact (contrib_pred k s' z Hk Hc Hin). }
    intro Hn. destruct (Nat.lt_ge_cases s NS) as [Hl|Hg].
    + destruct (G (NS - s) s (le_n _) Hn) as [Hc|Hc]; [exact Hc | lia].
    + destruct (needf_step us n _ s Hn) as [H0|(k & z & Hk & _ & Hin)].
      * apply contrib_out. apply existsb_exists in H0 as (y & Hy & E). apply Nat.eqb_eq in E. subst. exact Hy.
      * exfalso. specialize (Hwf k Hk s z Hin). unfold NS in Hg. lia.
  - intro Hc. induction Hc as [s Hs|k s z Hk Hc IH Hin]; [apply need_nx; exact Hs|].
    exact (need_pred Hwf k s z Hk IH Hin).
Qed.

(* ---- the last layer that reads a signal ------------------------------------------------------------ *)
Definition lam (s : nat) : nat :=
  Nat.max (if existsb (Nat.eqb s) nx then S dmax else 0)
          (maxl (map (fun k => if need (n + k) && reads k s then dep (n + k) else 0) (seq 0 (length us)))).

Lemma lam_nx : forall s, In s nx -> S dmax <= lam s.
Proof.
  intros s H. unfold lam. replace (existsb (Nat.eqb s) nx) with true; [lia|].
  symmetry. apply existsb_exists. exists s. split; [exact H | apply Nat.eqb_refl].
Qed.

Lemma lam_reader : forall k s z, k < length us -> need (n + k) = true -> In (s, z) (tin (nth k us tu0)) ->
  dep (n + k) <= lam s.
Proof.
  intros k s z Hk Hn Hin. unfold lam. apply Nat.max_le_iff. right.
  replace (dep (n + k)) with (if need (n + k) && reads k s then dep (n + k) else 0).
  - apply maxl_ge. apply in_map_iff. exists k. split; [reflexivity | apply in_seq; lia].
  - rewrite Hn. replace (reads k s) with true; [reflexivity|]. symmetry. apply reads_in. exists z. exact Hin.
Qed.

Lemma lam_need : wfn -> forall s, 1 <= lam s -> need s = true.
Proof.
  intros Hwf s H. unfold lam in H. destruct (existsb (Nat.eqb s) nx) eqn:E.
  - apply existsb_exists in E as (y & Hy & E). apply Nat.eqb_eq in E. subst. apply need_nx. exact Hy.
  - cbn [Nat.max] in H. destruct (maxl_pos _ H) as (a & Ha & H1). apply in_map_iff in Ha as (k & Ea & Hk).
    apply in_seq in Hk. destruct (need (n + k) && reads k s) eqn:F; [|lia].
    apply andb_true_iff in F as [F1 F2]. apply reads_in in F2 as (z & Hin).
    exact (need_pred Hwf k s z ltac:(lia) F1 Hin).
Qed.

(* ---- the layers ----------------------------------------------------------------------------------- *)
Definition isunit (l s : nat) : bool := (n <=? s) && (dep s =? l).
Definition live (l s : nat) : bool := isunit l s && need s || (dep s <? l) && (l <? lam s).
Definition pos (l s : nat) : nat := match l with O => s | _ => length (filter (live l) (seq 0 s)) end.

Definition idr (p : nat) : tu := TU (-1) [(p, 1%Z)].
Definition reidx (l : nat) (u : tu) : tu := TU (tb u) (map (fun p => (pos l (fst p), snd p)) (tin u)).

(* the source of a row: a copy of a signal, or a unit *)
Inductive src := Cp (s : nat) | Un (s : nat).
Definition srcof (l s : nat) : src := if isunit l s then Un s else Cp s.
Definition srcs (l : nat) : list src := map (srcof l) (filter (live l) (seq 0 NS)).
Definition srcd : list src := map (fun i => srcof dmax (nth i nx 0)) (seq 0 n).
Definition mkrow (l : nat) (c : src) : tu :=
  match c with Cp s => idr (pos (l - 1) s) | Un s => reidx (l - 1) (nth (s - n) us tu0) end.

(* layers 1 .. d - 1 hold the live signals, layer d one row for each state bit *)
Definition lsrc : list (nat * list src) := map (fun l => (l, srcs l)) (seq 1 (dmax - 1)) ++ [(dmax, srcd)].
Definition lev : list (list tu) := map (fun p => map (mkrow (fst p)) (snd p)) lsrc.

(* ---- positions --------------------------------------------------------------------------------------- *)
Lemma seq_split : forall s M, s < M -> seq 0 M = seq 0 s ++ s :: seq (S s) (M - S s).
Proof. intros s M H. replace M with (s + S (M - S s)) at 1 by lia. rewrite seq_app. reflexivity. Qed.

Lemma nth_filter_count : forall (f : nat -> bool) s M d, s < M -> f s = true ->
  nth (length (filter f (seq 0 s))) (filter f (seq 0 M)) d = s.
Proof.
  intros f s M d Hs Hf. rewrite (seq_split s M Hs).
  rewrite filter_app. cbn [filter]. rewrite Hf. rewrite app_nth2 by lia. rewrite Nat.sub_diag. reflexivity.
Qed.

Lemma count_lt : forall (f : nat -> bool) s M, s < M -> f s = true ->
  length (filter f (seq 0 s)) < length (filter f (seq 0 M)).
Proof.
  intros f s M Hs Hf. rewrite (seq_split s M Hs).
  rewrite filter_app, length_app. cbn [filter]. rewrite Hf. cbn [length]. lia.
Qed.

Lemma count_mono : forall (f : nat -> bool) a b, a <= b -> length (filter f (seq 0 a)) <= length (filter f (seq 0 b)).
Proof.
  intros f a b H. replace b with (a + (b - a)) by lia. rewrite seq_app, filter_app, length_app. lia.
Qed.

Lemma pos_inj : forall l s1 s2, 1 <= l -> live l s1 = true -> live l s2 = true -> pos l s1 = pos l s2 -> s1 = s2.
Proof.
  intros l s1 s2 Hl H1 H2 E. destruct l as [|l]; [lia|]. cbn [pos] in E.
  destruct (Nat.lt_trichotomy s1 s2) as [Hlt|[Heq|Hgt]]; [|exact Heq|].
  - pose proof (count_lt (live (S l)) s1 s2 Hlt H1). lia.
  - pose proof (count_lt (live (S l)) s2 s1 Hgt H2). lia.
Qed.

(* ---- firing the rows ------------------------------------------------------------------------------ *)
Lemma wsum_reidx : forall l ins E,
  wsum (map (fun p => (pos l (fst p), snd p)) ins) E = wsum ins (fun i => E (pos l i)).
Proof. intros l ins E. induction ins as [|[i z] t IH]; cbn [map wsum fst snd]; [reflexivity|]. rewrite IH. reflexivity. Qed.

Lemma fire_reidx : forall l u E, fire (reidx l u) E = fire u (fun i => E (pos l i)).
Proof. intros l u E. unfold fire, reidx. cbn [tb tin]. rewrite wsum_reidx. reflexivity. Qed.

Lemma fire_idr : forall p E, fire (idr p) E = E p.
Proof. intros p E. unfold fire, idr. cbn [tb tin wsum]. destruct (E p); reflexivity. Qed.

Lemma lfire_nth : forall L y r, r < length L -> nth r (lfire L y) false = fire (nth r L tu0) (fun i => nth i y false).
Proof.
  intros L y r Hr. unfold lfire. rewrite (nth_indep _ false (fire tu0 (fun i => nth i y false))) by (rewrite length_map; exact Hr).
  exact (map_nth (fun u => fire u (fun i => nth i y false)) L tu0 r).
Qed.

(* a signal has a row in layer l, l <= d - 1 *)
Definition at_layer (l s : nat) : Prop := s < NS /\ (match l with O => s < n | _ => live l s = true end).

(* the predecessors of a unit of depth l + 1 that contributes, and a signal live at l + 1 or a next value of depth
   at most l, have rows in layer l *)
Lemma pred_at : wfn -> forall l k p z, k < length us -> need (n + k) = true -> dep (n + k) = S l ->
  In (p, z) (tin (nth k us tu0)) -> at_layer l p.
Proof.
  intros Hwf l k p z Hk Hn Hdk Hin. pose proof (dep_pred Hwf k p z Hk Hin) as Hlt.
  assert (Hp : p < NS) by (pose proof (Hwf k Hk p z Hin); unfold NS; lia).
  split; [exact Hp|]. destruct l as [|l].
  - apply (dep_zero Hwf p Hp). lia.
  - unfold live. pose proof (lam_reader k p z Hk Hn Hin) as Hlam.
    destruct (Nat.eq_dec (dep p) (S l)) as [E|E].
    + apply orb_true_iff. left. apply andb_true_iff. split; [|exact (need_pred Hwf k p z Hk Hn Hin)].
      unfold isunit. apply andb_true_iff. split; [|apply Nat.eqb_eq; exact E].
      apply Nat.leb_le. destruct (Nat.lt_ge_cases p n) as [Hl|Hg]; [rewrite dep_state in E by exact Hl; lia | exact Hg].
    + apply orb_true_iff. right. apply andb_true_iff. split; apply Nat.ltb_lt; lia.
Qed.

Lemma copy_at : wfn -> forall l s, s < NS -> dep s <= l -> S l < lam s -> at_layer l s.
Proof.
  intros Hwf l s Hs Hdl Hlam. split; [exact Hs|]. destruct l as [|l].
  - apply (dep_zero Hwf s Hs). lia.
  - unfold live. destruct (Nat.eq_dec (dep s) (S l)) as [E|E].
    + apply orb_true_iff. left. apply andb_true_iff. split; [|apply lam_need; [exact Hwf | lia]].
      unfold isunit. apply andb_true_iff. split; [|apply Nat.eqb_eq; exact E].
      apply Nat.leb_le. destruct (Nat.lt_ge_cases s n) as [Hl|Hg]; [rewrite dep_state in E by exact Hl; lia | exact Hg].
    + apply orb_true_iff. right. apply andb_true_iff. split; apply Nat.ltb_lt; lia.
Qed.

Section Sem.
Hypothesis Hwf : wfn.
Hypothesis Hnx : Forall (fun s => s < NS) nx.
Hypothesis Hlen : length nx = n.
Hypothesis Hd : 1 <= dmax.

Variable x : list bool.
Hypothesis Hx : length x = n.

Let F : nat -> bool := run us n (fun i => nth i x false).

Lemma F_unit : forall k, k < length us -> F (n + k) = fire (nth k us tu0) F.
Proof. intros k Hk. unfold F. apply run_at; [exact Hk | exact Hwf]. Qed.

Lemma F_state : forall i, i < n -> F i = nth i x false.
Proof. intros i Hi. unfold F. apply run_prefix. exact Hi. Qed.

(* the outputs after the first l layers, 0 <= l <= d - 1 *)
Definition yl (l : nat) : list bool := lmap (map (fun l' => map (mkrow l') (srcs l')) (seq 1 l)) x.

(* firing a row from the outputs of the layer before *)
Lemma fire_row : forall l (y : list bool) c,
  (forall s, at_layer l s -> nth (pos l s) y false = F s) ->
  (forall s, c = Un s -> exists k, k < length us /\ s = n + k /\ need s = true /\ dep s = S l) ->
  (forall s, c = Cp s -> at_layer l s) ->
  fire (mkrow (S l) c) (fun i => nth i y false) = F (match c with Cp s => s | Un s => s end).
Proof.
  intros l y c Hy HU HC. destruct c as [s|s]; cbn [mkrow]; replace (S l - 1) with l by lia.
  - rewrite fire_idr. apply Hy. apply HC. reflexivity.
  - destruct (HU s eq_refl) as (k & Hk & -> & Hn & Hdk). replace (n + k - n) with k by lia.
    rewrite fire_reidx, F_unit by exact Hk. apply fire_ext. intros p z Hin. apply Hy. exact (pred_at Hwf l k p z Hk Hn Hdk Hin).
Qed.

Lemma yl_spec : forall l, l <= dmax - 1 ->
  length (yl l) = (match l with O => n | _ => length (filter (live l) (seq 0 NS)) end)
  /\ forall s, at_layer l s -> nth (pos l s) (yl l) false = F s.
Proof.
  induction l as [|l IH]; intros Hl.
  - unfold yl. cbn. split; [exact Hx|]. intros s [Hs Hn]. symmetry. apply F_state. exact Hn.
  - destruct (IH ltac:(lia)) as [IL IY].
    assert (E : yl (S l) = lfire (map (mkrow (S l)) (srcs (S l))) (yl l)).
    { unfold yl. rewrite seq_S, map_app. cbn [map]. replace (1 + l) with (S l) by lia. apply lmap_app. }
    rewrite E. unfold lfire, srcs. rewrite !length_map. split; [reflexivity|].
    intros s [Hs Hlv]. cbn [pos]. fold (lfire (map (mkrow (S l)) (map (srcof (S l)) (filter (live (S l)) (seq 0 NS)))) (yl l)).
    pose proof (count_lt (live (S l)) s NS Hs Hlv) as Hc.
    rewrite lfire_nth by (rewrite !length_map; exact Hc).
    rewrite nth_map_d with (d' := Cp 0) by (rewrite length_map; exact Hc).
    rewrite nth_map_d with (d' := 0) by exact Hc. rewrite nth_filter_count by assumption.
    assert (Hfs : F s = F (match srcof (S l) s with Cp s' => s' | Un s' => s' end))
      by (unfold srcof; destruct (isunit (S l) s); reflexivity).
    rewrite Hfs. apply fire_row; [exact IY | |].
    + intros s' Hs'. unfold srcof in Hs'. destruct (isunit (S l) s) eqn:Hu; [|discriminate]. injection Hs' as <-.
      unfold isunit in Hu. apply andb_true_iff in Hu as [H1 H2]. apply Nat.leb_le in H1. apply Nat.eqb_eq in H2.
      exists (s - n). split; [unfold NS in Hs; lia|]. split; [lia|]. split; [|exact H2].
      unfold live in Hlv. apply orb_true_iff in Hlv as [Hlv|Hlv].
      * apply andb_true_iff in Hlv as [_ Hlv]. exact Hlv.
      * apply andb_true_iff in Hlv as [Hlv _]. apply Nat.ltb_lt in Hlv. lia.
    + intros s' Hs'. unfold srcof in Hs'. destruct (isunit (S l) s) eqn:Hu; [discriminate|]. injection Hs' as <-.
      unfold live in Hlv. rewrite Hu in Hlv. cbn [andb orb] in Hlv. apply andb_true_iff in Hlv as [H1 H2].
      apply Nat.ltb_lt in H1, H2. apply (copy_at Hwf); [exact Hs | lia | lia].
Qed.

Theorem lev_correct : lmap lev x = netstep n us nx x.
Proof.
  unfold lev, lsrc. rewrite map_app. cbn [map fst snd]. rewrite lmap_app.
  fold (yl (dmax - 1)). destruct (yl_spec (dmax - 1) (le_n _)) as [_ IY].
  unfold netstep. fold F. unfold lfire, srcd. rewrite !map_map.
  transitivity (map F (map (fun i => nth i nx 0) (seq 0 n))).
  2: { f_equal. rewrite <- Hlen. apply map_nth_self. }
  rewrite map_map. apply map_ext_in. intros i Hi. apply in_seq in Hi.
  set (s := nth i nx 0).
  assert (Hs : s < NS) by (rewrite Forall_forall in Hnx; apply Hnx; apply nth_In; lia).
  assert (Hin : In s nx) by (apply nth_In; lia).
  assert (Hfs : F s = F (match srcof dmax s with Cp s' => s' | Un s' => s' end))
    by (unfold srcof; destruct (isunit dmax s); reflexivity).
  rewrite Hfs. replace (mkrow dmax) with (mkrow (S (dmax - 1))) by (f_equal; lia).
  apply fire_row; [exact IY | |].
  - intros s' Hs'. unfold srcof in Hs'. destruct (isunit dmax s) eqn:Hu; [|discriminate]. injection Hs' as <-.
    unfold isunit in Hu. apply andb_true_iff in Hu as [H1 H2]. apply Nat.leb_le in H1. apply Nat.eqb_eq in H2.
    exists (s - n). split; [unfold NS in Hs; lia|]. split; [lia|]. split; [apply need_nx; exact Hin | lia].
  - intros s' Hs'. unfold srcof in Hs'. destruct (isunit dmax s) eqn:Hu; [discriminate|]. injection Hs' as <-.
    pose proof (lam_nx s Hin). pose proof (dep_le s Hs).
    assert (dep s <> dmax \/ s < n) as Hne.
    { unfold isunit in Hu. apply andb_false_iff in Hu as [Hu|Hu]; [right; apply Nat.leb_gt in Hu; exact Hu|].
      left. apply Nat.eqb_neq in Hu. lia. }
    apply (copy_at Hwf); [exact Hs | | lia].
    destruct Hne as [Hne|Hne]; [lia | rewrite dep_state by exact Hne; lia].
Qed.

End Sem.

(* ---- the shape of the layered map ----------------------------------------------------------------------- *)
Lemma length_lev : 1 <= dmax -> length lev = dmax.
Proof. intro H. unfold lev, lsrc. rewrite length_map, length_app, length_map, length_seq. cbn [length]. lia. Qed.

(* the width of layer l, 0 <= l <= d - 1 *)
Definition width (l : nat) : nat := match l with O => n | _ => length (filter (live l) (seq 0 NS)) end.

Lemma pos_lt : forall l s, at_layer l s -> pos l s < width l.
Proof.
  intros l s [Hs Hl]. destruct l as [|l]; cbn [pos width]; [exact Hl|]. apply count_lt; assumption.
Qed.

Lemma pos_inj_at : forall l s1 s2, at_layer l s1 -> at_layer l s2 -> pos l s1 = pos l s2 -> s1 = s2.
Proof.
  intros l s1 s2 [_ H1] [_ H2] E. destruct l as [|l]; [exact E|]. apply (pos_inj (S l)); [lia | exact H1 | exact H2 | exact E].
Qed.

Lemma nodup_map_inj : forall (f : nat -> nat) (P : nat -> Prop) l, NoDup l -> (forall a, In a l -> P a) ->
  (forall a b, P a -> P b -> f a = f b -> a = b) -> NoDup (map f l).
Proof.
  intros f P l Hnd. induction Hnd as [|a t Ha Hnd IH]; intros HP Hinj; cbn [map]; constructor.
  - intro Hin. apply in_map_iff in Hin as (b & Eb & Hb). apply Ha.
    rewrite (Hinj b a (HP b (or_intror Hb)) (HP a (or_introl eq_refl)) Eb) in Hb. exact Hb.
  - apply IH; [intros b Hb; apply HP; right; exact Hb | exact Hinj].
Qed.

(* a row built from a source of layer l + 1 reads distinct positions of layer l with weights -1 or 1 *)
Lemma mkrow_ok : wfn -> Forall (fun u => NoDup (map fst (tin u)) /\ forall i z, In (i, z) (tin u) -> z = 1%Z \/ z = (-1)%Z) us ->
  forall l c, (forall s, c = Un s -> exists k, k < length us /\ s = n + k /\ need s = true /\ dep s = S l) ->
  (forall s, c = Cp s -> at_layer l s) -> rowok (width l) (mkrow (S l) c).
Proof.
  intros Hwf Hus l c HU HC. destruct c as [s|s]; cbn [mkrow]; replace (S l - 1) with l by lia.
  - specialize (HC s eq_refl). unfold rowok, idr. cbn [tin map fst]. split; [constructor; [intros []|constructor]|].
    intros i z [E|[]]. injection E as <- <-. split; [apply pos_lt; exact HC | left; reflexivity].
  - destruct (HU s eq_refl) as (k & Hk & -> & Hn & Hdk). replace (n + k - n) with k by lia.
    rewrite Forall_forall in Hus. destruct (Hus (nth k us tu0) (nth_In _ _ Hk)) as [Hnd Hz].
    unfold rowok, reidx. cbn [tin]. split.
    + rewrite map_map. cbn [fst]. rewrite <- (map_map fst (pos l)).
      apply (nodup_map_inj (pos l) (at_layer l) (map fst (tin (nth k us tu0)))); [exact Hnd | |].
      * intros p Hp. apply in_map_iff in Hp as ([p' z] & <- & Hin). exact (pred_at Hwf l k p' z Hk Hn Hdk Hin).
      * intros a b Ha Hb E. exact (pos_inj_at l a b Ha Hb E).
    + intros i z Hin. apply in_map_iff in Hin as ([p z'] & E & Hin). cbn [fst snd] in E. injection E as <- <-.
      split; [apply pos_lt; exact (pred_at Hwf l k p z' Hk Hn Hdk Hin) | exact (Hz p z' Hin)].
Qed.

Lemma srcs_spec : forall l c, In c (srcs l) ->
  (forall s, c = Un s -> exists k, k < length us /\ s = n + k /\ need s = true /\ dep s = l)
  /\ (forall s, c = Cp s -> s < NS /\ dep s < l < lam s).
Proof.
  intros l c Hc. unfold srcs in Hc. apply in_map_iff in Hc as (s & <- & Hs). apply filter_In in Hs as [Hs Hlv].
  apply in_seq in Hs. unfold srcof. destruct (isunit l s) eqn:Hu; split; intros s' E; try discriminate; injection E as <-.
  - unfold isunit in Hu. apply andb_true_iff in Hu as [H1 H2]. apply Nat.leb_le in H1. apply Nat.eqb_eq in H2.
    exists (s - n). split; [unfold NS in Hs; lia|]. split; [lia|]. split; [|exact H2].
    unfold live in Hlv. apply orb_true_iff in Hlv as [Hlv|Hlv].
    + apply andb_true_iff in Hlv as [_ Hlv]. exact Hlv.
    + apply andb_true_iff in Hlv as [Hlv _]. apply Nat.ltb_lt in Hlv. lia.
  - unfold live in Hlv. rewrite Hu in Hlv. cbn [andb orb] in Hlv. apply andb_true_iff in Hlv as [H1 H2].
    apply Nat.ltb_lt in H1, H2. split; [lia | lia].
Qed.

Lemma srcd_spec : Forall (fun s => s < NS) nx -> length nx = n -> 1 <= dmax -> forall c, In c srcd ->
  (forall s, c = Un s -> exists k, k < length us /\ s = n + k /\ need s = true /\ dep s = dmax)
  /\ (forall s, c = Cp s -> s < NS /\ dep s < dmax /\ S dmax <= lam s).
Proof.
  intros Hnx Hlen Hd c Hc. unfold srcd in Hc. apply in_map_iff in Hc as (i & <- & Hi). apply in_seq in Hi.
  set (s := nth i nx 0). assert (Hs : s < NS) by (rewrite Forall_forall in Hnx; apply Hnx; apply nth_In; lia).
  assert (Hin : In s nx) by (apply nth_In; lia).
  unfold srcof. destruct (isunit dmax s) eqn:Hu; split; intros s' E; try discriminate; injection E as <-.
  - unfold isunit in Hu. apply andb_true_iff in Hu as [H1 H2]. apply Nat.leb_le in H1. apply Nat.eqb_eq in H2.
    exists (s - n). split; [unfold NS in Hs; lia|]. split; [lia|]. split; [apply need_nx; exact Hin | exact H2].
  - split; [exact Hs|]. split; [|apply lam_nx; exact Hin]. pose proof (dep_le s Hs).
    unfold isunit in Hu. apply andb_false_iff in Hu as [Hu|Hu].
    + apply Nat.leb_gt in Hu. rewrite dep_state by exact Hu. lia.
    + apply Nat.eqb_neq in Hu. lia.
Qed.

(* Lemma 2.5: the layered map has depth d, every row reads distinct rows of the layer before with weights -1 or 1,
   and it computes the next state of the netlist *)
Theorem lev_ok : wfn -> Forall (fun s => s < NS) nx -> length nx = n -> 1 <= dmax ->
  Forall (fun u => NoDup (map fst (tin u)) /\ forall i z, In (i, z) (tin u) -> z = 1%Z \/ z = (-1)%Z) us ->
  length lev = dmax /\ lwf n lev /\ length (last lev []) = n
  /\ forall x, length x = n -> lmap lev x = netstep n us nx x.
Proof.
  intros Hwf Hnx Hlen Hd Hus. split; [apply length_lev; exact Hd|]. split; [|split].
  - (* layer l + 1 reads layer l *)
    assert (G : forall m, m <= dmax - 1 ->
      lwf (width m) (map (fun p => map (mkrow (fst p)) (snd p))
                       (map (fun l => (l, srcs l)) (seq (S m) (dmax - 1 - m)) ++ [(dmax, srcd)]))).
    { intros m Hm. remember (dmax - 1 - m) as c eqn:Ec. revert m Hm Ec. induction c as [|c IH]; intros m Hm Ec.
      - cbn [seq map app lwf fst snd]. split; [|exact I].
        apply Forall_forall. intros r Hr. apply in_map_iff in Hr as (c & <- & Hc).
        destruct (srcd_spec Hnx Hlen Hd c Hc) as [HU HC].
        replace dmax with (S m) by lia. apply mkrow_ok; [exact Hwf | exact Hus | |].
        + intros s E. destruct (HU s E) as (k & Hk & Es & Hn & Hdk). exists k. repeat split; try assumption. lia.
        + intros s E. destruct (HC s E) as (Hs & Hdl & Hlam). apply copy_at; [exact Hwf | exact Hs | lia | lia].
      - cbn [seq map app lwf fst snd]. split.
        + apply Forall_forall. intros r Hr. apply in_map_iff in Hr as (c' & <- & Hc).
          destruct (srcs_spec (S m) c' Hc) as [HU HC]. apply mkrow_ok; [exact Hwf | exact Hus | exact HU |].
          intros s E. destruct (HC s E) as (Hs & Hdl & Hlam). apply copy_at; [exact Hwf | exact Hs | lia | lia].
        + rewrite length_map. unfold srcs. rewrite length_map.
          replace (length (filter (live (S m)) (seq 0 NS))) with (width (S m)) by reflexivity.
          apply IH; lia. }
    specialize (G 0 ltac:(lia)). unfold lev, lsrc. replace (dmax - 1) with (dmax - 1 - 0) by lia. exact G.
  - unfold lev, lsrc. rewrite map_app. cbn [map fst snd]. rewrite last_last. unfold srcd. rewrite !length_map, length_seq.
    reflexivity.
  - intros x Hx. apply lev_correct; assumption.
Qed.

(* every row is an identity unit H(x - 1) or a unit of the netlist with its bias and weights *)
Theorem lev_rows : forall l cs c, In (l, cs) lsrc -> In c cs ->
  match c with
  | Cp s => exists p, mkrow l c = idr p
  | Un s => tb (mkrow l c) = tb (nth (s - n) us tu0) /\ map snd (tin (mkrow l c)) = map snd (tin (nth (s - n) us tu0))
  end.
Proof.
  intros l cs c _ _. destruct c as [s|s]; cbn [mkrow].
  - eexists. reflexivity.
  - unfold reidx. cbn [tb tin]. split; [reflexivity|]. rewrite map_map. reflexivity.
Qed.

(* the occurrences of unit s among the rows *)
Definition isun (s : nat) (c : src) : bool := match c with Un s' => Nat.eqb s' s | Cp _ => false end.

Lemma length_filter_map : forall (A B : Type) (p : B -> bool) (f : A -> B) (L : list A),
  length (filter p (map f L)) = length (filter (fun a => p (f a)) L).
Proof. intros A B p f L. induction L as [|a t IH]; cbn [map filter]; [reflexivity|]. destruct (p (f a)); cbn [length]; lia. Qed.

Lemma length_filter_filter : forall (A : Type) (p q : A -> bool) (L : list A),
  length (filter p (filter q L)) = length (filter (fun a => q a && p a) L).
Proof.
  intros A p q L. induction L as [|a t IH]; [reflexivity|]. cbn [filter].
  destruct (q a); cbn [andb filter]; [destruct (p a); cbn [length]; lia | exact IH].
Qed.

Lemma filter_eqb_seq : forall (g : nat -> bool) M s, s < M ->
  length (filter (fun s' => g s' && Nat.eqb s' s) (seq 0 M)) = if g s then 1 else 0.
Proof.
  intros g M s Hs. rewrite (seq_split s M Hs), filter_app, length_app. cbn [filter]. rewrite Nat.eqb_refl, andb_true_r.
  assert (Z : forall a k, (a > s \/ a + k <= s) -> length (filter (fun s' => g s' && Nat.eqb s' s) (seq a k)) = 0).
  { intros a k Hak. revert a Hak. induction k as [|k IH]; intros a Hak; [reflexivity|]. cbn [seq filter].
    destruct (Nat.eqb_spec a s); [lia|]. rewrite andb_false_r. apply IH. lia. }
  rewrite (Z 0 s) by lia. destruct (g s); cbn [length]; rewrite (Z (S s) (M - S s)) by lia; reflexivity.
Qed.

Lemma count_srcs : forall l s, s < NS -> length (filter (isun s) (srcs l)) = if isunit l s && need s then 1 else 0.
Proof.
  intros l s Hs. unfold srcs. rewrite length_filter_map, length_filter_filter.
  rewrite (filter_ext _ (fun a => (live l a && isunit l a) && Nat.eqb a s)).
  - rewrite filter_eqb_seq by exact Hs. destruct (isunit l s) eqn:Hu.
    + unfold live. rewrite Hu. unfold isunit in Hu. apply andb_true_iff in Hu as [_ Hu]. apply Nat.eqb_eq in Hu.
      rewrite Hu, Nat.ltb_irrefl. destruct (need s); reflexivity.
    + rewrite andb_false_r. reflexivity.
  - intro a. unfold srcof. destruct (isunit l a); cbn [isun]; [rewrite andb_true_r; reflexivity|].
    rewrite !andb_false_r. reflexivity.
Qed.

Lemma count_map_srcof : forall d s (L : list nat),
  length (filter (fun a => isun s (srcof d a)) L) = if isunit d s then length (filter (fun a => Nat.eqb a s) L) else 0.
Proof.
  intros d s L. induction L as [|a t IH]; [destruct (isunit d s); reflexivity|]. cbn [filter].
  destruct (Nat.eqb_spec a s) as [->|E].
  - assert (Hs : isun s (srcof d s) = isunit d s)
      by (unfold srcof; destruct (isunit d s); cbn [isun]; [apply Nat.eqb_refl | reflexivity]).
    rewrite Hs. destruct (isunit d s) eqn:Hu.
    + cbn [length]. rewrite IH. reflexivity.
    + exact IH.
  - assert (Hf : isun s (srcof d a) = false)
      by (unfold srcof; destruct (isunit d a); cbn [isun]; [apply Nat.eqb_neq; exact E | reflexivity]).
    rewrite Hf. exact IH.
Qed.

Lemma count_srcd : length nx = n -> forall s,
  length (filter (isun s) srcd) = if isunit dmax s then length (filter (fun a => Nat.eqb a s) nx) else 0.
Proof.
  intros Hlen s. unfold srcd. rewrite <- (map_map (fun i => nth i nx 0) (srcof dmax)).
  rewrite <- Hlen, map_nth_self, length_filter_map. apply count_map_srcof.
Qed.

Lemma length_filter_concat : forall (A : Type) (p : A -> bool) (L : list (list A)),
  length (filter p (concat L)) = fold_right Nat.add 0 (map (fun l => length (filter p l)) L).
Proof.
  intros A p L. induction L as [|l t IH]; [reflexivity|]. cbn [concat map fold_right]. rewrite filter_app, length_app, IH.
  reflexivity.
Qed.

Lemma sum_indicator : forall x a k, fold_right Nat.add 0 (map (fun l => if Nat.eqb x l then 1 else 0) (seq a k))
  = if (a <=? x) && (x <? a + k) then 1 else 0.
Proof.
  intros x a k. revert a. induction k as [|k IH]; intro a.
  - cbn [seq map fold_right]. destruct (Nat.leb_spec a x), (Nat.ltb_spec x (a + 0)); cbn [andb]; lia.
  - cbn [seq map fold_right]. rewrite IH. destruct (Nat.eqb_spec x a) as [->|E].
    + rewrite Nat.leb_refl. destruct (Nat.leb_spec (S a) a); [lia|]. cbn [andb]. destruct (Nat.ltb_spec a (a + S k)); lia.
    + destruct (Nat.leb_spec (S a) x), (Nat.ltb_spec x (S a + k)), (Nat.leb_spec a x), (Nat.ltb_spec x (a + S k));
        cbn [andb]; lia.
Qed.

Lemma sum_zero : forall (L : list nat) (f : nat -> nat), (forall l, In l L -> f l = 0) ->
  fold_right Nat.add 0 (map f L) = 0.
Proof.
  intros L f H. induction L as [|a t IH]; [reflexivity|]. cbn [map fold_right].
  rewrite (H a (or_introl eq_refl)), IH by (intros l Hl; apply H; right; exact Hl). reflexivity.
Qed.

Lemma filter_none : forall (L : list nat) s, (forall a, In a L -> a <> s) -> length (filter (fun a => Nat.eqb a s) L) = 0.
Proof.
  intros L s H. induction L as [|a t IH]; [reflexivity|]. cbn [filter].
  destruct (Nat.eqb_spec a s) as [E|E]; [exfalso; exact (H a (or_introl eq_refl) E)|].
  apply IH. intros b Hb. apply H. right. exact Hb.
Qed.

(* each unit that contributes occurs once if its depth is less than d and once for each state bit whose next value it
   is if its depth is d; no other unit occurs *)
Theorem lev_units : wfn -> length nx = n -> forall s, n <= s < NS ->
  length (filter (isun s) (concat (map snd lsrc)))
  = if need s then (if dep s <? dmax then 1 else length (filter (fun a => Nat.eqb a s) nx)) else 0.
Proof.
  intros Hwf Hlen s Hs. unfold lsrc. rewrite map_app, concat_app, filter_app, length_app, map_map. cbn [map snd concat].
  rewrite app_nil_r, count_srcd by exact Hlen. rewrite length_filter_concat, map_map.
  pose proof (dep_unit_pos Hwf (s - n) ltac:(unfold NS in Hs; lia)) as Hpos. replace (n + (s - n)) with s in Hpos by lia.
  pose proof (dep_le s (proj2 Hs)) as Hle.
  assert (Hu : forall l, isunit l s = Nat.eqb (dep s) l)
    by (intro l; unfold isunit; destruct (Nat.leb_spec n s); [reflexivity | lia]).
  rewrite Hu. destruct (need s) eqn:Hn.
  - rewrite (map_ext_in _ (fun l => if Nat.eqb (dep s) l then 1 else 0)).
    2: { intros l Hl. rewrite count_srcs by lia. rewrite Hu, Hn, andb_true_r. reflexivity. }
    rewrite sum_indicator. destruct (Nat.ltb_spec (dep s) dmax) as [Hlt|Hge].
    + replace ((1 <=? dep s) && (dep s <? 1 + (dmax - 1))) with true
        by (symmetry; apply andb_true_iff; split; [apply Nat.leb_le | apply Nat.ltb_lt]; lia).
      destruct (Nat.eqb_spec (dep s) dmax); [lia | reflexivity].
    + replace ((1 <=? dep s) && (dep s <? 1 + (dmax - 1))) with false
        by (symmetry; apply andb_false_iff; right; apply Nat.ltb_ge; lia).
      destruct (Nat.eqb_spec (dep s) dmax); [reflexivity | lia].
  - rewrite sum_zero by (intros l _; rewrite count_srcs by lia; rewrite Hn, andb_false_r; reflexivity).
    destruct (Nat.eqb (dep s) dmax); [|reflexivity]. rewrite filter_none; [reflexivity|].
    intros a Ha E. subst a. rewrite (need_nx s Ha) in Hn. discriminate.
Qed.

(* the number of rows of each layer *)
Theorem lev_widths : length nx = n -> map (@length tu) lev = map width (seq 1 (dmax - 1)) ++ [n].
Proof.
  intro Hlen. unfold lev, lsrc. rewrite map_app, map_app, !map_map. cbn [map fst snd]. f_equal.
  - apply map_ext_in. intros l Hl. apply in_seq in Hl. unfold srcs. rewrite !length_map.
    destruct l as [|l]; [lia | reflexivity].
  - unfold srcd. rewrite !length_map, length_seq. reflexivity.
Qed.

(* the size of the layered map: its rows are those of the layers 1 .. d - 1 and one for each state bit *)
Theorem lev_size : length nx = n ->
  fold_right Nat.add 0 (map (@length tu) lev) = fold_right Nat.add 0 (map width (seq 1 (dmax - 1))) + n.
Proof.
  intro Hlen. unfold lev, lsrc. rewrite map_app, map_app, fold_right_app, !map_map. cbn [map fold_right fst snd].
  unfold srcd. rewrite !length_map, length_seq.
  assert (G : forall (L : list nat) c, (forall l, In l L -> 1 <= l) ->
            fold_right Nat.add c (map (fun l => length (map (mkrow l) (srcs l))) L) = fold_right Nat.add 0 (map width L) + c).
  { intros L c HL. induction L as [|a t IH]; [reflexivity|]. cbn [map fold_right].
    rewrite IH by (intros l Hl; apply HL; right; exact Hl). unfold srcs. rewrite !length_map.
    destruct a as [|a]; [specialize (HL 0 (or_introl eq_refl)); lia|]. cbn [width]. lia. }
  rewrite G by (intros l Hl; apply in_seq in Hl; lia). lia.
Qed.

End Lev.