(* The host netlist read in its own order, for every A >= 3: every weight is 1 or -1, no unit reads the
   unit of a next value, and a unit that reads a state bit comes no later than the unit of that bit's next
   value. So the unit of each next value may write its state bit in place. *)
From Stdlib Require Import List Bool Arith ZArith Lia.
Import ListNotations.
From TH Require Import TwoHead Net Gates Host HostWF PhysFacts.

(* ---- flat indices and block coordinates ------------------------------------------------------------ *)
Lemma flat_loc : forall bs, Forall shaped bs -> forall k, k < boff bs (length bs) ->
  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.
Proof.
  intros bs Hsh. 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.
Qed.

Lemma idx_inj : forall bs n x y, valid bs n x -> valid bs n y -> idx n bs x = idx n bs y -> x = y.
Proof.
  intros bs n [i|b j r] [i'|b' j' r'] Hx Hy E; cbn [valid idx] in Hx, Hy, E.
  - subst. reflexivity.
  - lia.
  - lia.
  - destruct Hx as (Hb & Hj & Hr). destruct Hy as (Hb' & Hj' & Hr').
    assert (Bb : b = b').
    { destruct (Nat.lt_trichotomy b b') as [L|[L|L]]; [|exact L|].
      - pose proof (boff_mono bs b' b L ltac:(lia)). nia.
      - pose proof (boff_mono bs b b' L ltac:(lia)). nia. }
    subst b'. assert (Jj : j = j') by nia. subst j'. f_equal. lia.
Qed.

(* ---- the inputs of the last four blocks ------------------------------------------------------------ *)
Ltac solve_disj := repeat (first [left; reflexivity | right]); reflexivity.

Ltac lits3 Hin := let H := fresh "H" in
  first [pose proof (in_sin_lits _ _ _ (or_introl Hin)) as H | pose proof (in_sin_lits _ _ _ (or_intror Hin)) as H];
  cbn [map fst] in H; destruct H as [<-|[<-|[<-|[]]]]; solve_disj.

Ltac lits2 Hin := let H := fresh "H" in
  first [pose proof (in_sin_lits _ _ _ (or_introl Hin)) as H | pose proof (in_sin_lits _ _ _ (or_intror Hin)) as H];
  cbn [map fst] in H; destruct H as [<-|[<-|[]]]; solve_disj.

Section Last.
Variable A : nat.
Hypothesis HA : 3 <= A.

Lemma len_hb : length (hblocks A) = tb A + 15.
Proof. apply length_hblocks. lia. Qed.

Lemma inputs_NP : forall j r x w, j < A -> r < 1 -> In (x, w) (sin (sunit (hblocks A) (tb A + 11) j r)) ->
  x = SRef 1 j 0 \/ x = SRef (tb A + 9) j 0 \/ x = SRef (tb A + 10) j 0.
Proof.
  intros j r x w Hj Hr Hin. unfold sunit in Hin. rewrite nth_hb_post in Hin by lia.
  cbn [nth hpost bshape blk_NP] in Hin. destruct r as [|r]; [|lia]. cbn [nth] in Hin. lits3 Hin.
Qed.

Lemma inputs_HAB : forall j r x w, j < 1 -> r < 3 -> In (x, w) (sin (sunit (hblocks A) (tb A + 12) j r)) ->
  x = haltb A \/ x = SRef (tb A + 8) 0 0 \/ x = SRef 9 0 0 \/ x = SRef 0 (SW A - 4) 0
  \/ x = SRef (tb A + 12) 0 0 \/ x = SRef (tb A + 12) 0 1.
Proof.
  intros j r x w Hj Hr Hin. unfold sunit in Hin. rewrite nth_hb_post in Hin by lia.
  cbn [nth hpost bshape blk_HAB] in Hin.
  assert (Hr3 : r = 0 \/ r = 1 \/ r = 2) by lia.
  destruct Hr3 as [->|[->| ->]]; cbn [nth] in Hin; lits3 Hin.
Qed.

Lemma inputs_WR : forall j r x w, r < 1 -> In (x, w) (sin (sunit (hblocks A) (tb A + 13) j r)) ->
  x = rbit A (A - 1 - j mod A) \/ x = SRef 10 (j / A) 0.
Proof.
  intros j r x w Hr Hin. unfold sunit in Hin. rewrite nth_hb_post in Hin by lia.
  cbn [nth hpost bshape blk_WR] in Hin. destruct r as [|r]; [|lia]. cbn [nth] in Hin. lits2 Hin.
Qed.

Lemma inputs_NM : forall j r x w, r < 1 -> In (x, w) (sin (sunit (hblocks A) (tb A + 14) j r)) ->
  x = SRef 13 j 0 \/ x = SRef (tb A + 13) j 0.
Proof.
  intros j r x w Hr Hin. unfold sunit in Hin. rewrite nth_hb_post in Hin by lia.
  cbn [nth hpost bshape blk_NM] in Hin. destruct r as [|r]; [|lia]. cbn [nth] in Hin. lits2 Hin.
Qed.

Lemma rbit_not : forall k b j r, rbit A k = SRef b j r -> b < tb A + 8.
Proof.
  intros k b j r. unfold rbit. destruct (Nat.eqb_spec k 0); intro E; injection E as <- _ _; unfold tb; lia.
Qed.

Lemma bsize_post : forall i, i < 15 -> bsize (nth (tb A + i) (hblocks A) blk0) = bsize (nth i (hpost A) blk0).
Proof. intros i Hi. rewrite nth_hb_post by lia. reflexivity. Qed.

Lemma bcount_post : forall i, i < 15 -> bcount (nth (tb A + i) (hblocks A) blk0) = bcount (nth i (hpost A) blk0).
Proof. intros i Hi. rewrite nth_hb_post by lia. reflexivity. Qed.

(* the units of the next values are valid units *)
Lemma hout_valid : forall i, i < nH A -> valid (hblocks A) (nH A) (hout A i).
Proof.
  intros i Hi. pose proof (SW_ge A HA). unfold hout.
  destruct (Nat.ltb_spec i A); [|destruct (Nat.eqb_spec i A)]; cbn [valid]; rewrite len_hb;
    (split; [lia|]); rewrite ?bcount_post, ?bsize_post by lia;
    cbn [nth hpost bcount bsize blk_NP blk_HAB blk_NM]; unfold nH in Hi; nia.
Qed.

(* no unit reads the unit of a next value *)
Theorem reads_no_out : forall b j r x w,
  b < length (hblocks A) -> j < bcount (nth b (hblocks A) blk0) -> r < bsize (nth b (hblocks A) blk0) ->
  In (x, w) (sin (sunit (hblocks A) b j r)) -> forall i, i < nH A -> x <> hout A i.
Proof.
  intros b j r x w Hb Hj Hr Hin i Hi E. rewrite len_hb in Hb.
  pose proof (proj2 (host_wf A HA) b j r x w ltac:(rewrite len_hb; lia) Hj Hr Hin) as Hrb.
  unfold hout in E.
  destruct (Nat.ltb_spec i A) as [Hia|Hia]; [|destruct (Nat.eqb_spec i A) as [Hie|Hie]]; subst x;
    cbn [reads_back] in Hrb.
  - (* SRef (tb + 11) i 0 *)
    destruct Hrb as [(Hb' & _ & _)|(Eb & Ej & Hr')].
    + assert (Hc : b = tb A + 12 \/ b = tb A + 13 \/ b = tb A + 14) by lia.
      destruct Hc as [->|[->| ->]].
      * rewrite bcount_post in Hj by lia. rewrite bsize_post in Hr by lia.
        cbn [nth hpost bcount bsize blk_HAB] in Hj, Hr.
        destruct (inputs_HAB j r _ w Hj Hr Hin) as [E|[E|[E|[E|[E|E]]]]]; try discriminate;
          injection E; unfold tb; lia.
      * rewrite bsize_post in Hr by lia. cbn [nth hpost bsize blk_WR] in Hr.
        destruct (inputs_WR j r _ w Hr Hin) as [E|E]; [|injection E; unfold tb; lia].
        pose proof (rbit_not (A - 1 - j mod A) _ _ _ (eq_sym E)). lia.
      * rewrite bsize_post in Hr by lia. cbn [nth hpost bsize blk_NM] in Hr.
        destruct (inputs_NM j r _ w Hr Hin) as [E|E]; injection E; unfold tb; lia.
    + subst b. rewrite bsize_post in Hr by lia. cbn [nth hpost bsize blk_NP] in Hr. lia.
  - (* SRef (tb + 12) 0 2 *)
    destruct Hrb as [(Hb' & _ & _)|(Eb & Ej & Hr')].
    + assert (Hc : b = tb A + 13 \/ b = tb A + 14) by lia.
      destruct Hc as [->| ->].
      * rewrite bsize_post in Hr by lia. cbn [nth hpost bsize blk_WR] in Hr.
        destruct (inputs_WR j r _ w Hr Hin) as [E|E]; [|injection E; unfold tb; lia].
        pose proof (rbit_not (A - 1 - j mod A) _ _ _ (eq_sym E)). lia.
      * rewrite bsize_post in Hr by lia. cbn [nth hpost bsize blk_NM] in Hr.
        destruct (inputs_NM j r _ w Hr Hin) as [E|E]; injection E; unfold tb; lia.
    + subst b. rewrite bsize_post in Hr by lia. cbn [nth hpost bsize blk_HAB] in Hr. lia.
  - (* SRef (tb + 14) (i - A - 1) 0 *)
    destruct Hrb as [(Hb' & _ & _)|(Eb & Ej & Hr')]; [lia|].
    subst b. rewrite bsize_post in Hr by lia. cbn [nth hpost bsize blk_NM] in Hr. lia.
Qed.

(* the state bits the last four blocks read: the halt bit, read by the halt latch alone *)
Lemma last_reads_halt : forall b j r i w, tb A + 11 <= b < tb A + 15 ->
  j < bcount (nth b (hblocks A) blk0) -> r < bsize (nth b (hblocks A) blk0) ->
  In (SIn i, w) (sin (sunit (hblocks A) b j r)) -> b = tb A + 12 /\ i = A.
Proof.
  intros b j r i w Hb Hj Hr Hin.
  assert (Hc : b = tb A + 11 \/ b = tb A + 12 \/ b = tb A + 13 \/ b = tb A + 14) by lia.
  destruct Hc as [->|[->|[->| ->]]].
  - rewrite bcount_post in Hj by lia. rewrite bsize_post in Hr by lia.
    cbn [nth hpost bcount bsize blk_NP] in Hj, Hr.
    destruct (inputs_NP j r _ w Hj Hr Hin) as [E|[E|E]]; discriminate.
  - rewrite bcount_post in Hj by lia. rewrite bsize_post in Hr by lia.
    cbn [nth hpost bcount bsize blk_HAB] in Hj, Hr.
    destruct (inputs_HAB j r _ w Hj Hr Hin) as [E|[E|[E|[E|[E|E]]]]]; try discriminate.
    unfold haltb in E. injection E as <-. split; reflexivity.
  - rewrite bsize_post in Hr by lia. cbn [nth hpost bsize blk_WR] in Hr.
    destruct (inputs_WR j r _ w Hr Hin) as [E|E]; [|discriminate].
    unfold rbit in E. destruct (Nat.eqb (A - 1 - j mod A) 0); discriminate.
  - rewrite bsize_post in Hr by lia. cbn [nth hpost bsize blk_NM] in Hr.
    destruct (inputs_NM j r _ w Hr Hin) as [E|E]; discriminate.
Qed.

(* a unit that reads state bit i comes no later than the unit of its next value *)
Theorem reads_in_before : forall b j r i w,
  b < length (hblocks A) -> j < bcount (nth b (hblocks A) blk0) -> r < bsize (nth b (hblocks A) blk0) ->
  In (SIn i, w) (sin (sunit (hblocks A) b j r)) ->
  i < nH A /\ idx (nH A) (hblocks A) (SRef b j r) <= idx (nH A) (hblocks A) (hout A i).
Proof.
  intros b j r i w Hb Hj Hr Hin.
  pose proof (proj2 (host_wf A HA) b j r (SIn i) w Hb Hj Hr Hin) as Hi. cbn [reads_back] in Hi.
  split; [exact Hi|]. rewrite len_hb in Hb.
  pose proof (hout_valid i Hi) as Hv.
  destruct (Nat.lt_ge_cases b (tb A + 11)) as [Hlt|Hge].
  - (* every unit before the next counter comes before every unit of a next value *)
    assert (Hb2 : forall b' j' r', tb A + 11 <= b' -> b' < tb A + 15 ->
                  idx (nH A) (hblocks A) (SRef b j r) < idx (nH A) (hblocks A) (SRef b' j' r')).
    { intros b' j' r' H1 H2. cbn [idx].
      pose proof (boff_mono (hblocks A) b' b ltac:(lia) ltac:(rewrite len_hb; lia)). nia. }
    unfold hout in *. destruct (Nat.ltb i A); [|destruct (Nat.eqb i A)]; apply Nat.lt_le_incl, Hb2; lia.
  - destruct (last_reads_halt b j r i w ltac:(lia) Hj Hr Hin) as [-> ->].
    unfold hout. rewrite Nat.ltb_irrefl, Nat.eqb_refl. cbn [idx].
    rewrite bcount_post in Hj by lia. rewrite bsize_post in Hr by lia.
    cbn [nth hpost bcount bsize blk_HAB] in Hj, Hr.
    rewrite bsize_post by lia. cbn [nth hpost bsize blk_HAB]. lia.
Qed.

(* every weight of the host is 1 or -1 *)
Lemma ter_hblocks : forall b j, Forall ter (bshape (nth b (hblocks A) blk0) j).
Proof.
  intros b j.
  assert (Hm : forall f n, (forall i, ter (f i)) -> Forall ter (map f (seq 0 n)))
    by (intros f n Hf; apply Forall_map, Forall_forall; intros; apply Hf).
  assert (Hsu : forall bb ws, Forall (fun p => snd p = 1%Z \/ snd p = (-1)%Z) ws -> ter (SU bb ws)).
  { intros bb ws H x w0 Hin. rewrite Forall_forall in H. exact (H (x, w0) Hin). }
  destruct (Nat.lt_ge_cases b (length (hblocks A))) as [Hb|Hb];
    [|rewrite nth_overflow by exact Hb; cbn [bshape blk0]; constructor].
  rewrite len_hb in Hb.
  destruct (Nat.lt_ge_cases b 17) as [Hp|Hp].
  { rewrite nth_hb_pre by exact Hp.
    assert (Hc : b = 0 \/ b = 1 \/ b = 2 \/ b = 3 \/ b = 4 \/ b = 5 \/ b = 6 \/ b = 7 \/ b = 8 \/ b = 9
                 \/ b = 10 \/ b = 11 \/ b = 12 \/ b = 13 \/ b = 14 \/ b = 15 \/ b = 16) by lia.
    destruct Hc as [->|[->|[->|[->|[->|[->|[->|[->|[->|[->|[->|[->|[->|[->|[->|[->| ->]]]]]]]]]]]]]]]];
      cbn [nth hpre bshape blk_P blk_HP blk_FA blk_Q blk_F blk_PA blk_PB blk_AC blk_W
           blk_RX blk_RY blk_KP blk_X blk_Y blk_PR];
      apply Forall_forall; intros u Hu; destruct Hu as [<-|[]]; first [apply ter_and | apply ter_or]. }
  destruct (Nat.lt_ge_cases b (tb A)) as [HT|HT].
  { unfold tb in HT. replace b with (16 + (b - 16)) by lia.
    rewrite nth_hb_T by lia. cbn [bshape blk_T].
    constructor; [apply ter_and | apply Forall_map, Forall_forall; intros; apply ter_and]. }
  replace b with (tb A + (b - tb A)) by lia. rewrite nth_hb_post by lia.
  assert (Hi : b - tb A < 15) by lia. remember (b - tb A) as i eqn:Ei. clear Ei.
  assert (Hc : i = 0 \/ i = 1 \/ i = 2 \/ i = 3 \/ i = 4 \/ i = 5 \/ i = 6 \/ i = 7 \/ i = 8 \/ i = 9
               \/ i = 10 \/ i = 11 \/ i = 12 \/ i = 13 \/ i = 14) by lia.
  assert (Hw : forall bb ws, (forall p, In p ws -> snd p = 1%Z \/ snd p = (-1)%Z) -> ter (SU bb ws))
    by (intros bb ws H; apply Hsu, Forall_forall, H).
  destruct Hc as [->|[->|[->|[->|[->|[->|[->|[->|[->|[->|[->|[->|[->|[->| ->]]]]]]]]]]]]]];
    cbn [nth hpost bshape blk_C blk_AX blk_NX blk_XN blk_EQ blk_R0 blk_H blk_R blk_LEQ
         blk_TA blk_TB blk_NP blk_HAB blk_WR blk_NM map];
    apply Forall_forall; intros u Hu;
    repeat (destruct Hu as [<-|Hu];
            [first [apply ter_and | apply ter_or
                   | apply Hw; intros p Hq; repeat (destruct Hq as [<-|Hq]; [cbn [snd]; lia|]); destruct Hq] |]);
    destruct Hu.
Qed.

End Last.

(* ---- the same facts on the flat netlist ------------------------------------------------------------- *)
Section Flat.
Variable A : nat.
Hypothesis HA : 3 <= A.

Let us := host_units A.

Lemma hb_shaped : Forall shaped (hblocks A).
Proof. exact (proj1 (host_wf A HA)). Qed.

Lemma length_hu : length us = boff (hblocks A) (length (hblocks A)).
Proof. apply length_compile. exact hb_shaped. Qed.

Lemma hu_loc : forall k, k < length us -> exists b j r,
  b < length (hblocks A) /\ j < bcount (nth b (hblocks A) blk0) /\ r < bsize (nth b (hblocks A) blk0)
  /\ nH A + k = idx (nH A) (hblocks A) (SRef b j r)
  /\ nth k us tu0 = flat_unit (nH A) (hblocks A) (sunit (hblocks A) b j r).
Proof.
  intros k Hk. rewrite length_hu in Hk.
  destruct (flat_loc (hblocks A) hb_shaped k Hk) as (b & j & r & Hb & Hj & Hr & ->).
  exists b, j, r. split; [exact Hb|]. split; [exact Hj|]. split; [exact Hr|]. split.
  - cbn [idx]. lia.
  - apply compile_nth; [exact hb_shaped | exact Hb | exact Hj | exact Hr].
Qed.

(* an input of a unit, in block coordinates *)
Lemma hu_input : forall k i w, k < length us -> In (i, w) (tin (nth k us tu0)) ->
  exists b j r x, b < length (hblocks A) /\ j < bcount (nth b (hblocks A) blk0) /\ r < bsize (nth b (hblocks A) blk0)
  /\ nH A + k = idx (nH A) (hblocks A) (SRef b j r) /\ In (x, w) (sin (sunit (hblocks A) b j r))
  /\ i = idx (nH A) (hblocks A) x /\ valid (hblocks A) (nH A) x.
Proof.
  intros k i w Hk Hin. destruct (hu_loc k Hk) as (b & j & r & Hb & Hj & Hr & Hidx & Hu).
  rewrite Hu in Hin. unfold flat_unit in Hin. cbn [tin] in Hin.
  apply in_map_iff in Hin as ([x w'] & E & Hx). cbn [fst snd] in E. injection E as <- <-.
  exists b, j, r, x. refine (conj Hb (conj Hj (conj Hr (conj Hidx (conj Hx (conj eq_refl _)))))).
  pose proof (proj2 (host_wf A HA) b j r x w' Hb Hj Hr Hx) as Hrb.
  destruct x as [i|b' j' r']; cbn [reads_back valid] in *; [exact Hrb|].
  destruct Hrb as [(H1 & H2 & H3)|(-> & -> & H3)]; [split; [lia | auto] | split; [exact Hb | split; [exact Hj | lia]]].
Qed.

Lemma hnext_nth : forall i, i < nH A -> nth i (host_next A) 0 = idx (nH A) (hblocks A) (hout A i).
Proof. intros i Hi. unfold host_next. apply nth_map_seq. exact Hi. Qed.

Lemma hout_inj : forall i i', hout A i = hout A i' -> i = i'.
Proof.
  intros i i'. unfold hout.
  destruct (Nat.ltb_spec i A), (Nat.ltb_spec i' A); try destruct (Nat.eqb_spec i A); try destruct (Nat.eqb_spec i' A);
    intro E; try injection E; intros; unfold tb in *; lia.
Qed.

Theorem host_weights : forall k i w, k < length us -> In (i, w) (tin (nth k us tu0)) -> w = 1%Z \/ w = (-1)%Z.
Proof.
  intros k i w Hk Hin. destruct (hu_input k i w Hk Hin) as (b & j & r & x & _ & _ & _ & _ & Hx & _).
  unfold sunit in Hx. pose proof (ter_hblocks A HA b j) as T.
  destruct (Nat.lt_ge_cases r (length (bshape (nth b (hblocks A) blk0) j))) as [Hr|Hr].
  - exact (ter_nth _ r T x w Hx).
  - rewrite nth_overflow in Hx by exact Hr. destruct Hx.
Qed.

Theorem host_no_out : forall k i w, k < length us -> In (i, w) (tin (nth k us tu0)) ->
  forall i', i' < nH A -> i <> nth i' (host_next A) 0.
Proof.
  intros k i w Hk Hin i' Hi' E. destruct (hu_input k i w Hk Hin) as (b & j & r & x & Hb & Hj & Hr & _ & Hx & -> & Hv).
  rewrite hnext_nth in E by exact Hi'.
  apply (idx_inj _ _ _ _ Hv (hout_valid A HA i' Hi')) in E.
  exact (reads_no_out A HA b j r x w Hb Hj Hr Hx i' Hi' E).
Qed.

Theorem host_in_before : forall k i w, k < length us -> In (i, w) (tin (nth k us tu0)) -> i < nH A ->
  nH A + k <= nth i (host_next A) 0.
Proof.
  intros k i w Hk Hin Hi. destruct (hu_input k i w Hk Hin) as (b & j & r & x & Hb & Hj & Hr & Hidx & Hx & Ex & Hv).
  destruct x as [i0|b' j' r']; cbn [idx] in Ex; [subst i0|lia].
  rewrite hnext_nth by exact Hi. rewrite Hidx.
  exact (proj2 (reads_in_before A HA b j r i w Hb Hj Hr Hx)).
Qed.

Lemma host_next_range : forall i, i < nH A -> nH A <= nth i (host_next A) 0 < nH A + length us.
Proof.
  intros i Hi. rewrite hnext_nth by exact Hi. pose proof (hout_valid A HA i Hi) as Hv.
  pose proof (idx_lt_total _ _ _ Hv) as H. rewrite <- length_hu in H.
  unfold hout in *. destruct (Nat.ltb i A); [|destruct (Nat.eqb i A)]; cbn [idx] in *;
    (split; [nia | exact H]).
Qed.

Lemma host_next_inj : forall i i', i < nH A -> i' < nH A -> nth i (host_next A) 0 = nth i' (host_next A) 0 -> i = i'.
Proof.
  intros i i' Hi Hi' E. rewrite !hnext_nth in E by assumption.
  apply hout_inj. exact (idx_inj _ _ _ _ (hout_valid A HA i Hi) (hout_valid A HA i' Hi') E).
Qed.

Lemma host_reads_below : forall k, k < length us -> reads_below (nth k us tu0) (nH A + k).
Proof. intros k Hk. apply wf_reads_below; [exact (host_wf A HA) | exact Hk]. Qed.

Lemma length_host_next : length (host_next A) = nH A.
Proof. unfold host_next. rewrite length_map, length_seq. reflexivity. Qed.

End Flat.