(* Netlists of ternary threshold units, flat and in blocks, and the valuations they compute. *)
From Stdlib Require Import List Bool Arith ZArith Lia.
Import ListNotations.
From TH Require Import TwoHead.

(* ---- flat netlists ------------------------------------------------------------- *)
(* a unit: its bias and its predecessors, each a signal index with a weight *)
Record tu := TU { tb : Z; tin : list (nat * Z) }.

Fixpoint wsum (ins : list (nat * Z)) (E : nat -> bool) : Z :=
  match ins with [] => 0%Z | (i, w) :: t => (w * b2z (E i) + wsum t E)%Z end.

Definition fire (u : tu) (E : nat -> bool) : bool := Z.leb 0 (tb u + wsum (tin u) E).

Definition setv (E : nat -> bool) (m : nat) (b : bool) : nat -> bool :=
  fun i => if Nat.eqb i m then b else E i.

(* the units evaluated in order, the first at signal m *)
Fixpoint run (us : list tu) (m : nat) (E : nat -> bool) : nat -> bool :=
  match us with [] => E | u :: t => run t (S m) (setv E m (fire u E)) end.

(* one step of a clocked netlist on n state bits: next value i is signal nx_i *)
Definition netstep (n : nat) (us : list tu) (nx : list nat) (x : list bool) : list bool :=
  let F := run us n (fun i => nth i x false) in map F nx.

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

Lemma wsum_ext : forall ins E E', (forall i w, In (i, w) ins -> E i = E' i) -> wsum ins E = wsum ins E'.
Proof.
  induction ins as [|[i w] t IH]; intros E E' H; cbn [wsum]; [reflexivity|].
  rewrite (H i w (or_introl eq_refl)), (IH E E'); [reflexivity|].
  intros i' w' Hin. apply (H i' w'). right. exact Hin.
Qed.

Lemma fire_ext : forall u E E', (forall i w, In (i, w) (tin u) -> E i = E' i) -> fire u E = fire u E'.
Proof. intros u E E' H. unfold fire. rewrite (wsum_ext _ E E' H). reflexivity. Qed.

Definition reads_below (u : tu) (m : nat) : Prop := forall i w, In (i, w) (tin u) -> i < m.

Definition tu0 : tu := TU 0 [].

Lemma run_at : forall us m E j, j < length us ->
  (forall k, k < length us -> reads_below (nth k us tu0) (m + k)) ->
  run us m E (m + j) = fire (nth j us tu0) (run us m E).
Proof.
  induction us as [|u t IH]; intros m E j Hj Hwf; cbn [length] in Hj, Hwf; [lia|].
  destruct j as [|j].
  - cbn [nth run]. rewrite Nat.add_0_r, run_prefix by lia. unfold setv at 1. rewrite Nat.eqb_refl.
    apply fire_ext. intros i w Hin. rewrite run_prefix by (specialize (Hwf 0 ltac:(lia) i w Hin); lia).
    unfold setv. destruct (Nat.eqb_spec i m); [specialize (Hwf 0 ltac:(lia) i w Hin); lia|reflexivity].
  - cbn [nth run]. replace (m + S j) with (S m + j) by lia. apply IH; [lia|].
    intros k Hk i w Hin. specialize (Hwf (S k) ltac:(lia) i w Hin). lia.
Qed.

(* ---- netlists in blocks ---------------------------------------------------------- *)
(* a signal is a state bit, or unit r of member j of block b *)
Inductive sref := SIn (i : nat) | SRef (b j r : nat).
Record su := SU { sb : Z; sin : list (sref * Z) }.
(* a block: count members of size units each; shape j lists the units of member j *)
Record block := Blk { bcount : nat; bsize : nat; bshape : nat -> list su }.

Definition blk0 : block := Blk 0 0 (fun _ => []).
Definition su0 : su := SU 0 [].

Fixpoint boff (bs : list block) (b : nat) : nat :=
  match bs, b with
  | _, O => 0
  | [], S _ => 0
  | blk :: t, S b' => bcount blk * bsize blk + boff t b'
  end.

Definition idx (n : nat) (bs : list block) (x : sref) : nat :=
  match x with
  | SIn i => i
  | SRef b j r => n + boff bs b + bsize (nth b bs blk0) * j + r
  end.

Definition flat_unit (n : nat) (bs : list block) (u : su) : tu :=
  TU (sb u) (map (fun p => (idx n bs (fst p), snd p)) (sin u)).

Definition flat_block (n : nat) (bs : list block) (blk : block) : list tu :=
  concat (map (fun j => map (flat_unit n bs) (bshape blk j)) (seq 0 (bcount blk))).

Definition compile (n : nat) (bs : list block) : list tu := concat (map (flat_block n bs) bs).

Definition shaped (blk : block) : Prop := forall j, j < bcount blk -> length (bshape blk j) = bsize blk.

Definition sunit (bs : list block) (b j r : nat) : su := nth r (bshape (nth b bs blk0) j) su0.

Definition valid (bs : list block) (n : nat) (x : sref) : Prop :=
  match x with
  | SIn i => i < n
  | SRef b j r => b < length bs /\ j < bcount (nth b bs blk0) /\ r < bsize (nth b bs blk0)
  end.

(* each unit reads state bits, units of earlier blocks, or earlier units of its own member *)
Definition reads_back (bs : list block) (n b j r : nat) (x : sref) : Prop :=
  match x with
  | SIn i => i < n
  | SRef b' j' r' => (b' < b /\ j' < bcount (nth b' bs blk0) /\ r' < bsize (nth b' bs blk0))
                     \/ (b' = b /\ j' = j /\ r' < r)
  end.

Definition wf (bs : list block) (n : nat) : Prop :=
  Forall shaped bs
  /\ forall b j r x w, b < length bs -> j < bcount (nth b bs blk0) -> r < bsize (nth b bs blk0) ->
       In (x, w) (sin (sunit bs b j r)) -> reads_back bs n b j r x.

Lemma length_concat_const : forall (A : Type) (ls : list (list A)) sz,
  Forall (fun l => length l = sz) ls -> length (concat ls) = length ls * sz.
Proof.
  intros A ls sz H. induction H as [|l ls Hl H IH]; [reflexivity|].
  cbn [concat length]. rewrite length_app, IH, Hl. lia.
Qed.

Lemma length_flat_block : forall n bs blk, shaped blk ->
  length (flat_block n bs blk) = bcount blk * bsize blk.
Proof.
  intros n bs blk Hs. unfold flat_block.
  rewrite length_concat_const with (sz := bsize blk).
  - rewrite length_map, length_seq. reflexivity.
  - apply Forall_forall. intros l Hl. apply in_map_iff in Hl as (j & <- & Hj).
    apply in_seq in Hj. rewrite length_map. apply Hs. lia.
Qed.

Lemma nth_concat_const : forall (A : Type) (d : A) (ls : list (list A)) sz j r,
  Forall (fun l => length l = sz) ls -> j < length ls -> r < sz ->
  nth (sz * j + r) (concat ls) d = nth r (nth j ls []) d.
Proof.
  intros A d ls sz j r Hall. revert j. induction Hall as [|l ls Hl Hall IH]; intros j Hj Hr;
    cbn [length] in Hj; [lia|].
  cbn [concat]. destruct j as [|j].
  - rewrite Nat.mul_0_r, Nat.add_0_l, app_nth1 by lia. reflexivity.
  - rewrite app_nth2 by nia. rewrite Hl. replace (sz * S j + r - sz) with (sz * j + r) by nia.
    apply IH; [lia | exact Hr].
Qed.

Lemma nth_map_seq : forall (A : Type) (f : nat -> A) c j d, j < c -> nth j (map f (seq 0 c)) d = f j.
Proof.
  intros A f c j d Hj. rewrite (nth_indep _ d (f 0)) by (rewrite length_map, length_seq; exact Hj).
  rewrite map_nth, seq_nth by exact Hj. reflexivity.
Qed.

Lemma nth_map_d : forall (A B : Type) (g : A -> B) l r d d', r < length l -> nth r (map g l) d = g (nth r l d').
Proof.
  intros A B g l r d d' Hr. rewrite (nth_indep _ d (g d')) by (rewrite length_map; exact Hr).
  apply map_nth.
Qed.

Lemma flat_block_nth : forall n bs blk j r, shaped blk -> j < bcount blk -> r < bsize blk ->
  nth (bsize blk * j + r) (flat_block n bs blk) tu0 = flat_unit n bs (nth r (bshape blk j) su0).
Proof.
  intros n bs blk j r Hs Hj Hr. unfold flat_block.
  rewrite nth_concat_const with (sz := bsize blk); [| | rewrite length_map, length_seq; exact Hj | exact Hr].
  - rewrite (nth_map_seq _ (fun j => map (flat_unit n bs) (bshape blk j)) (bcount blk) j [] Hj).
    apply nth_map_d. rewrite Hs by exact Hj. exact Hr.
  - apply Forall_forall. intros l Hl. apply in_map_iff in Hl as (j' & <- & Hj').
    apply in_seq in Hj'. rewrite length_map. apply Hs. lia.
Qed.

Lemma compile_nth : forall n bs b j r, Forall shaped bs -> b < length bs ->
  j < bcount (nth b bs blk0) -> r < bsize (nth b bs blk0) ->
  nth (boff bs b + bsize (nth b bs blk0) * j + r) (compile n bs) tu0
  = flat_unit n bs (sunit bs b j r).
Proof.
  intros n bs. unfold compile, sunit.
  assert (G : forall pre, Forall shaped pre ->
    forall b j r, b < length pre -> j < bcount (nth b pre blk0) -> r < bsize (nth b pre blk0) ->
    nth (boff pre b + bsize (nth b pre blk0) * j + r) (concat (map (flat_block n bs) pre)) tu0
    = flat_unit n bs (nth r (bshape (nth b pre blk0) j) su0)).
  { induction pre as [|blk t IH]; intros Hsh b j r Hb Hj Hr; cbn [length] in Hb; [lia|].
    inversion Hsh as [|? ? Hs1 Hst]; subst.
    cbn [map concat]. destruct b as [|b].
    - cbn [boff nth] in *. rewrite Nat.add_0_l, app_nth1.
      + apply flat_block_nth; assumption.
      + rewrite length_flat_block by exact Hs1. nia.
    - cbn [boff nth] in *. rewrite app_nth2 by (rewrite length_flat_block by exact Hs1; lia).
      rewrite length_flat_block by exact Hs1.
      replace (bcount blk * bsize blk + boff t b + bsize (nth b t blk0) * j + r - bcount blk * bsize blk)
        with (boff t b + bsize (nth b t blk0) * j + r) by lia.
      apply IH; [exact Hst | lia | exact Hj | exact Hr]. }
  intros b j r Hsh Hb Hj Hr. apply G; assumption.
Qed.

Lemma length_compile_gen : forall n bs0 bs, Forall shaped bs ->
  length (concat (map (flat_block n bs0) bs)) = boff bs (length bs).
Proof.
  intros n bs0 bs H. induction H as [|blk t Hs Ht IH]; [reflexivity|].
  cbn [map concat length boff]. rewrite length_app, length_flat_block by exact Hs. rewrite IH. reflexivity.
Qed.

Lemma length_compile : forall n bs, Forall shaped bs -> length (compile n bs) = boff bs (length bs).
Proof. intros n bs H. apply length_compile_gen. exact H. Qed.

Lemma boff_mono : forall bs b b', b' < b -> b <= length bs ->
  boff bs b' + bcount (nth b' bs blk0) * bsize (nth b' bs blk0) <= boff bs b.
Proof.
  induction bs as [|blk t IH]; intros b b' Hb Hl; cbn [length] in Hl; [lia|].
  destruct b as [|b]; [lia|]. destruct b' as [|b'].
  - cbn [boff nth]. lia.
  - cbn [boff nth]. specialize (IH b b' ltac:(lia) ltac:(lia)). lia.
Qed.

Lemma idx_lt : forall bs n b j r, b < length bs -> j < bcount (nth b bs blk0) -> r < bsize (nth b bs blk0) ->
  idx n bs (SRef b j r) < n + boff bs (S b).
Proof.
  intros bs n b j r Hb Hj Hr. cbn [idx].
  pose proof (boff_mono bs (S b) b ltac:(lia) ltac:(lia)). nia.
Qed.

Lemma idx_lt_total : forall bs n x, valid bs n x -> idx n bs x < n + boff bs (length bs).
Proof.
  intros bs n [i|b j r] Hv; cbn [valid] in Hv; [cbn [idx]; lia|].
  destruct Hv as (Hb & Hj & Hr). pose proof (idx_lt bs n b j r Hb Hj Hr).
  destruct (Nat.eq_dec (S b) (length bs)) as [E|E]; [rewrite <- E; exact H|].
  pose proof (boff_mono bs (length bs) (S b) ltac:(lia) ltac:(lia)). lia.
Qed.

(* the structured unit fired on a valuation of its signals *)
Fixpoint swsum (ins : list (sref * Z)) (V : sref -> bool) : Z :=
  match ins with [] => 0%Z | (x, w) :: t => (w * b2z (V x) + swsum t V)%Z end.

Definition sfire' (u : su) (V : sref -> bool) : bool := Z.leb 0 (sb u + swsum (sin u) V).

Lemma wsum_flat : forall n bs ins E,
  wsum (map (fun p => (idx n bs (fst p), snd p)) ins) E = swsum ins (fun x => E (idx n bs x)).
Proof. intros n bs ins E. induction ins as [|[x w] t IH]; cbn [map wsum swsum fst snd]; [reflexivity|]. rewrite IH. reflexivity. Qed.

Lemma swsum_ext : forall ins V V', (forall x w, In (x, w) ins -> V x = V' x) -> swsum ins V = swsum ins V'.
Proof.
  induction ins as [|[x w] t IH]; intros V V' H; cbn [swsum]; [reflexivity|].
  rewrite (H x w (or_introl eq_refl)), (IH V V'); [reflexivity|].
  intros x' w' Hin. apply (H x' w'). right. exact Hin.
Qed.

Lemma wf_reads_below : forall bs n, wf bs n ->
  forall k, k < length (compile n bs) -> reads_below (nth k (compile n bs) tu0) (n + k).
Proof.
  intros bs n [Hsh Hrd] k Hk. rewrite length_compile in Hk by exact Hsh.
  assert (Hloc : exists b j r, b < length bs /\ j < bcount (nth b bs blk0) /\ r < bsize (nth b bs blk0)
                               /\ k = boff bs b + bsize (nth b bs blk0) * j + r).
  { clear -Hk Hsh. revert k Hk. induction Hsh as [|blk t Hs Ht IH]; intros k Hk; cbn [boff length] in Hk; [lia|].
    destruct (Nat.lt_ge_cases k (bcount blk * bsize blk)) as [Hl|Hg].
    - exists 0, (k / bsize blk), (k mod bsize blk).
      assert (bsize blk <> 0) by (intro E; rewrite E in Hl; lia).
      cbn [nth length boff]. split; [lia|]. split; [apply Nat.Div0.div_lt_upper_bound; nia|].
      split; [apply Nat.mod_upper_bound; exact H|]. rewrite (Nat.div_mod k (bsize blk)) at 1 by exact H. lia.
    - destruct (IH (k - bcount blk * bsize blk) ltac:(lia)) as (b & j & r & Hb & Hj & Hr & Hk').
      exists (S b), j, r. cbn [nth length boff]. split; [lia|]. split; [exact Hj|]. split; [exact Hr|]. lia. }
  destruct Hloc as (b & j & r & Hb & Hj & Hr & ->).
  rewrite compile_nth by assumption.
  intros i w Hin. unfold flat_unit in Hin. cbn [tin] in Hin.
  apply in_map_iff in Hin as ([x w'] & Heq & Hin). cbn [fst snd] in Heq. injection Heq as <- <-.
  specialize (Hrd b j r x w' Hb Hj Hr Hin).
  destruct x as [i|b' j' r']; cbn [reads_back idx] in *; [lia|].
  destruct Hrd as [(Hb' & Hj' & Hr')|(-> & -> & Hr')]; [|lia].
  pose proof (boff_mono bs b b' Hb' ltac:(lia)). nia.
Qed.

(* the value of a state bit, and of a unit from the values of its predecessors *)
Lemma net_in : forall bs n X i, i < n -> run (compile n bs) n X (idx n bs (SIn i)) = X i.
Proof. intros bs n X i H. cbn [idx]. apply run_prefix. exact H. Qed.

Lemma net_unit : forall bs n X b j r, wf bs n ->
  b < length bs -> j < bcount (nth b bs blk0) -> r < bsize (nth b bs blk0) ->
  run (compile n bs) n X (idx n bs (SRef b j r))
  = sfire' (sunit bs b j r) (fun x => run (compile n bs) n X (idx n bs x)).
Proof.
  intros bs n X b j r Hwf Hb Hj Hr. cbn [idx].
  replace (n + boff bs b + bsize (nth b bs blk0) * j + r) with (n + (boff bs b + bsize (nth b bs blk0) * j + r)) by lia.
  rewrite run_at.
  - rewrite compile_nth by (destruct Hwf; assumption). unfold fire, flat_unit, sfire'. cbn [tb tin].
    rewrite wsum_flat. reflexivity.
  - rewrite length_compile by (destruct Hwf; assumption).
    pose proof (idx_lt_total bs n (SRef b j r) (conj Hb (conj Hj Hr))) as H. cbn [idx] in H. lia.
  - apply wf_reads_below. exact Hwf.
Qed.