(* Recipes (Section 6): the encoder describe of Lemma 6.7, as src/host.py computes it, with its decoding, its
   length bound and the case of equality; the optimality of the bound for strings with no two equal adjacent
   bytes (Proposition 6.6); and the completeness of the recipe language (Lemma 6.5). *)
From Stdlib Require Import List Bool Arith Lia.
Import ListNotations.
From TH Require Import Net Host HostWF HostSem Device Pstar.

Lemma nu_nu : forall b, b < 256 -> nu (nu b) = b.
Proof.
  intros b Hb. unfold nu. destruct (Nat.eq_dec b 0) as [->|E]; [reflexivity|].
  rewrite (Nat.mod_small (256 - b)) by lia. rewrite Nat.mod_small by lia. lia.
Qed.

Lemma nu_lt : forall b, nu b < 256.
Proof. intro b. unfold nu. apply Nat.mod_upper_bound. lia. Qed.

(* ---- the encoder ------------------------------------------------------------------------------------------ *)
(* the length of the run of bytes equal to x at the head of f *)
Fixpoint run_len (x : nat) (f : list nat) : nat :=
  match f with [] => 0 | y :: f' => if Nat.eqb y x then S (run_len x f') else 0 end.

(* four equal bytes begin at the head of f *)
Definition four_at (f : list nat) : bool :=
  match f with a :: b :: c :: d :: _ => Nat.eqb b a && Nat.eqb c a && Nat.eqb d a | _ => false end.

(* the bytes before the first position at which four equal bytes begin, at most cap of them *)
Fixpoint lit_scan (cap : nat) (f : list nat) : nat :=
  match cap with
  | O => 0
  | S c => match f with [] => 0 | _ :: f' => if four_at f then 0 else S (lit_scan c f') end
  end.

Fixpoint desc (fuel : nat) (f : list nat) : list token :=
  match fuel with
  | O => []
  | S fu =>
      match f with
      | [] => []
      | x :: _ =>
          let l := Nat.min (run_len x f) 127 in
          if Nat.leb 4 l then Rep (256 - l) (nu x) :: desc fu (skipn l f)
          else let k := lit_scan 127 f in Lit (map nu (firstn k f)) :: desc fu (skipn k f)
      end
  end.

Definition describe (f : list nat) : list token := desc (length f) f.

(* the bound of Lemmas 6.5 and 6.7: |f| + ceil(|f| / 127) + 1 *)
Definition bound (n : nat) : nat := n + (n + 126) / 127 + 1.

Definition nofour (f : list nat) : Prop := forall p, four_at (skipn p f) = false.

(* ---- runs and scans ------------------------------------------------------------------------------------- *)
Lemma firstn_run : forall x f l, l <= run_len x f -> firstn l f = repeat x l.
Proof.
  intros x f. induction f as [|y f IH]; intros l Hl; cbn [run_len] in Hl.
  - assert (l = 0) by lia. subst. reflexivity.
  - destruct l as [|l]; [reflexivity|]. destruct (Nat.eqb_spec y x) as [->|E]; [|lia].
    cbn [firstn repeat]. rewrite IH by lia. reflexivity.
Qed.

Lemma run_le : forall x f, run_len x f <= length f.
Proof. intros x f. induction f as [|y f IH]; cbn [run_len length]; [lia|]. destruct (Nat.eqb y x); lia. Qed.

Lemma four_run : forall x f, four_at (x :: f) = Nat.leb 4 (run_len x (x :: f)).
Proof.
  intros x [|b [|c [|d f]]]; cbn [four_at run_len]; rewrite Nat.eqb_refl; try reflexivity.
  - destruct (Nat.eqb b x); reflexivity.
  - destruct (Nat.eqb b x), (Nat.eqb c x); reflexivity.
  - destruct (Nat.eqb b x), (Nat.eqb c x), (Nat.eqb d x); reflexivity.
Qed.

Lemma lit_scan_le : forall cap f, lit_scan cap f <= cap /\ lit_scan cap f <= length f.
Proof.
  induction cap as [|c IH]; intros f; cbn [lit_scan]; [lia|]. destruct f as [|y f']; cbn [length]; [lia|].
  destruct (four_at (y :: f')); [lia|]. specialize (IH f'). lia.
Qed.

Lemma lit_scan_none : forall cap f p, p < lit_scan cap f -> four_at (skipn p f) = false.
Proof.
  induction cap as [|c IH]; intros f p Hp; cbn [lit_scan] in Hp; [lia|]. destruct f as [|y f']; [lia|].
  destruct (four_at (y :: f')) eqn:E; [lia|]. destruct p as [|p]; [exact E|]. cbn [skipn]. apply (IH f'). lia.
Qed.

Lemma lit_scan_stop : forall cap f, lit_scan cap f = cap \/ lit_scan cap f = length f
  \/ four_at (skipn (lit_scan cap f) f) = true.
Proof.
  induction cap as [|c IH]; intros f; cbn [lit_scan]; [left; reflexivity|]. destruct f as [|y f']; [right; left; reflexivity|].
  destruct (four_at (y :: f')) eqn:E; [right; right; exact E|]. cbn [skipn length].
  destruct (IH f') as [H|[H|H]]; [left; lia | right; left; lia | right; right; exact H].
Qed.

Lemma lit_scan_pos : forall cap f, 1 <= cap -> f <> [] -> four_at f = false -> 1 <= lit_scan cap f.
Proof. intros [|c] [|y f] Hc Hf E; [lia | lia | contradiction Hf; reflexivity | cbn [lit_scan]; rewrite E; lia]. Qed.

Lemma nofour_skipn : forall f k, nofour f -> nofour (skipn k f).
Proof. intros f k H p. rewrite skipn_skipn. apply H. Qed.

Lemma Forall_take : forall (T : Type) (P : T -> Prop) l k, Forall P l -> Forall P (firstn k l).
Proof. intros T P l k H. rewrite <- (firstn_skipn k l) in H. apply Forall_app in H. exact (proj1 H). Qed.

Lemma Forall_drop : forall (T : Type) (P : T -> Prop) l k, Forall P l -> Forall P (skipn k l).
Proof. intros T P l k H. rewrite <- (firstn_skipn k l) in H. apply Forall_app in H. exact (proj2 H). Qed.

(* ---- the recipe of a token list ---------------------------------------------------------------------------- *)
Lemma length_recipe_cons : forall k ks, length (recipe (k :: ks)) = length (enc k) + length (recipe ks).
Proof. intros k ks. unfold recipe. cbn [map concat]. rewrite <- app_assoc, length_app. reflexivity. Qed.

Lemma delta_cons : forall k ks, delta (k :: ks) = dec k ++ delta ks.
Proof. reflexivity. Qed.

Lemma ceil_add : forall m, (127 + m + 126) / 127 = 1 + (m + 126) / 127.
Proof. intro m. replace (127 + m + 126) with (1 * 127 + (m + 126)) by lia. rewrite Nat.div_add_l by lia. reflexivity. Qed.

Lemma ceil_mono : forall a b, a <= b -> (a + 126) / 127 <= (b + 126) / 127.
Proof. intros a b H. apply Nat.Div0.div_le_mono. lia. Qed.

Lemma ceil_small : forall n, 1 <= n -> n <= 127 -> (n + 126) / 127 = 1.
Proof. intros n H1 H2. symmetry. apply Nat.div_unique with (r := n + 126 - 127); lia. Qed.

(* ---- Lemma 6.7 --------------------------------------------------------------------------------------------- *)
Lemma desc_spec : forall fuel f, length f <= fuel -> Forall (fun x => x < 256) f ->
  let r := length (recipe (desc fuel f)) in
  delta (desc fuel f) = f /\ Forall tok_ok (desc fuel f) /\ r <= bound (length f)
  /\ (four_at f = true -> r + 2 <= bound (length f))
  /\ (nofour f -> r = bound (length f)) /\ (r = bound (length f) -> nofour f).
Proof.
  induction fuel as [|fu IH]; intros f Hl Hb r.
  - assert (f = []) by (destruct f; [reflexivity | cbn in Hl; lia]). subst f. unfold r. cbn.
    split; [reflexivity|]. split; [constructor|]. split; [lia|]. split; [discriminate|]. split; [reflexivity|].
    intros _ p. destruct p; reflexivity.
  - destruct f as [|x f0].
    + unfold r. cbn.
      split; [reflexivity|]. split; [constructor|]. split; [lia|]. split; [discriminate|]. split; [reflexivity|].
      intros _ p. destruct p; reflexivity.
    + assert (Hx : x < 256) by (inversion Hb; assumption).
      unfold r. cbn [desc]. cbv zeta.
      remember (x :: f0) as f eqn:Ef.
      assert (Hfl : length f = S (length f0)) by (subst f; reflexivity).
      set (l := Nat.min (run_len x f) 127).
      assert (Hrun : run_len x f >= 1) by (subst f; cbn [run_len]; rewrite Nat.eqb_refl; lia).
      pose proof (run_le x f) as Hrl.
      destruct (Nat.leb_spec 4 l) as [H4|H4].
      * (* a run of at least four: a repeat token *)
        set (f2 := skipn l f).
        assert (Hlf : length f = l + length f2) by (unfold f2; rewrite length_skipn; unfold l in *; lia).
        destruct (IH f2 ltac:(lia) ltac:(unfold f2; apply Forall_drop; exact Hb)) as (D & T & B & B2 & E1 & E2).
        cbv zeta in *. set (r2 := length (recipe (desc fu f2))) in *.
        assert (Hdec : dec (Rep (256 - l) (nu x)) = firstn l f).
        { cbn [dec]. rewrite nu_nu by exact Hx. replace (256 - (256 - l)) with l by (unfold l in *; lia).
          symmetry. apply firstn_run. unfold l. lia. }
        rewrite length_recipe_cons. cbn [enc length]. fold r2.
        assert (Hb2 : bound (length f2) + 4 <= bound (length f)).
        { unfold bound. rewrite Hlf. pose proof (ceil_mono (length f2) (l + length f2) ltac:(lia)). lia. }
        refine (conj _ (conj _ (conj _ (conj _ (conj _ _))))).
        -- rewrite delta_cons, D, Hdec. apply firstn_skipn.
        -- constructor; [|exact T]. cbn [tok_ok]. split; [unfold l in *; lia | apply nu_lt].
        -- lia.
        -- intros _. lia.
        -- intro Hn. exfalso. specialize (Hn 0). cbn [skipn] in Hn. subst f. rewrite four_run in Hn.
           apply Nat.leb_gt in Hn. unfold l in *. lia.
        -- intro E. lia.
      * (* a literal token *)
        assert (Hf4 : four_at f = false).
        { subst f. rewrite four_run. apply Nat.leb_gt. unfold l in *. lia. }
        set (k := lit_scan 127 f).
        destruct (lit_scan_le 127 f) as [Hk1 Hk2]. fold k in Hk1, Hk2.
        assert (Hk0 : 1 <= k) by (apply lit_scan_pos; [lia | subst f; discriminate | exact Hf4]).
        set (f2 := skipn k f).
        assert (Hlf : length f = k + length f2) by (unfold f2; rewrite length_skipn; lia).
        destruct (IH f2 ltac:(lia) ltac:(unfold f2; apply Forall_drop; exact Hb)) as (D & T & B & B2 & E1 & E2).
        cbv zeta in *. set (r2 := length (recipe (desc fu f2))) in *.
        assert (Hg : Forall (fun x => x < 256) (firstn k f)) by (apply Forall_take; exact Hb).
        rewrite length_recipe_cons. cbn [enc length]. rewrite length_map, length_firstn, Nat.min_l by lia. fold r2.
        assert (Hnone : forall p, p < k -> four_at (skipn p f) = false) by (intros p Hp; apply (lit_scan_none 127); exact Hp).
        assert (Hnf : nofour f2 -> nofour f).
        { intros Hn p. destruct (Nat.lt_ge_cases p k) as [Hp|Hp]; [apply Hnone; exact Hp|].
          replace p with ((p - k) + k) by lia. rewrite <- skipn_skipn. apply (Hn (p - k)). }
        refine (conj _ (conj _ (conj _ (conj _ (conj _ _))))).
        -- rewrite delta_cons, D. cbn [dec]. rewrite map_map.
           rewrite (map_ext_in _ (fun x => x)) by (intros y Hy; rewrite Forall_forall in Hg; apply nu_nu, Hg, Hy).
           rewrite map_id. apply firstn_skipn.
        -- constructor; [|exact T]. cbn [tok_ok]. rewrite length_map, length_firstn, Nat.min_l by lia.
           split; [lia|]. apply Forall_map, Forall_forall. intros y _. apply nu_lt.
        -- destruct (lit_scan_stop 127 f) as [S1|[S1|S1]]; fold k in S1.
           ++ assert (Hb' : bound (length f) = 128 + bound (length f2))
                by (unfold bound; rewrite Hlf, S1, ceil_add; lia). lia.
           ++ assert (Hf2 : length f2 = 0) by lia. rewrite Hf2 in B.
              assert (Hc : (length f + 126) / 127 = 1) by (apply ceil_small; lia).
              unfold bound in *. rewrite Hc. cbn in B. lia.
           ++ destruct (Nat.eq_dec k 127) as [K7|K7].
              ** assert (Hb' : bound (length f) = 128 + bound (length f2))
                   by (unfold bound; rewrite Hlf, K7, ceil_add; lia). lia.
              ** specialize (B2 S1). unfold bound in *. rewrite Hlf. pose proof (ceil_mono (length f2) (k + length f2) ltac:(lia)). lia.
        -- intro H. rewrite Hf4 in H. discriminate.
        -- intro Hn. assert (Hn2 : nofour f2) by (apply nofour_skipn; exact Hn).
           specialize (E1 Hn2). destruct (lit_scan_stop 127 f) as [S1|[S1|S1]]; fold k in S1.
           ++ assert (Hb' : bound (length f) = 128 + bound (length f2))
                by (unfold bound; rewrite Hlf, S1, ceil_add; lia). lia.
           ++ assert (Hf2 : length f2 = 0) by lia. rewrite Hf2 in E1.
              assert (Hc : (length f + 126) / 127 = 1) by (apply ceil_small; lia).
              unfold bound in *. rewrite Hc. cbn in E1. lia.
           ++ exfalso. specialize (Hn k). rewrite S1 in Hn. discriminate.
        -- intro E. destruct (lit_scan_stop 127 f) as [S1|[S1|S1]]; fold k in S1.
           ++ assert (Hb' : bound (length f) = 128 + bound (length f2))
                by (unfold bound; rewrite Hlf, S1, ceil_add; lia).
              apply Hnf, E2. lia.
           ++ apply Hnf. intro p. unfold f2. rewrite skipn_skipn, skipn_all2 by lia. reflexivity.
           ++ exfalso. specialize (B2 S1). unfold bound in *. rewrite Hlf in E.
              pose proof (ceil_mono (length f2) (k + length f2) ltac:(lia)). lia.
Qed.

Theorem describe_spec : forall f, Forall (fun x => x < 256) f ->
  delta (describe f) = f /\ Forall tok_ok (describe f) /\ length (recipe (describe f)) <= bound (length f)
  /\ (length (recipe (describe f)) = bound (length f) <-> nofour f).
Proof.
  intros f Hb. destruct (desc_spec (length f) f (le_n _) Hb) as (D & T & B & _ & E1 & E2).
  unfold describe. refine (conj D (conj T (conj B (conj E2 E1)))).
Qed.

(* ---- the steps of P* on describe f (Theorem 7.5) ------------------------------------------------------------- *)
(* the steps of Proposition 6.13 on the recipe of ks *)
Definition psteps (ks : list token) : nat := dsteps ks + 7 + 1 + (6 * length (recipe ks) + 2).

Lemma psteps_nil : psteps [] = 16.
Proof. reflexivity. Qed.

Lemma psteps_cons : forall k ks, psteps (k :: ks) = tsteps k + 6 * length (enc k) + psteps ks.
Proof. intros k ks. unfold psteps. rewrite length_recipe_cons. cbn [dsteps fold_right]. fold (dsteps ks). lia. Qed.

Lemma desc_steps : forall fuel f, length f <= fuel -> Forall (fun x => x < 256) f ->
  2 * psteps (desc fuel f) <= 25 * length f + 58
  /\ (four_at f = true -> 2 * psteps (desc fuel f) <= 25 * length f + 32)
  /\ (nofour f -> psteps (desc fuel f) = 11 * length f + 13 * ((length f + 126) / 127) + 16).
Proof.
  induction fuel as [|fu IH]; intros f Hl Hb.
  - assert (f = []) by (destruct f; [reflexivity | cbn in Hl; lia]). subst f. cbn [desc length].
    rewrite psteps_nil. split; [lia|]. split; [discriminate|]. intros _. reflexivity.
  - destruct f as [|x f0].
    + cbn [desc length]. rewrite psteps_nil. split; [lia|]. split; [discriminate|]. intros _. reflexivity.
    + assert (Hx : x < 256) by (inversion Hb; assumption).
      cbn [desc]. cbv zeta.
      remember (x :: f0) as f eqn:Ef.
      set (l := Nat.min (run_len x f) 127).
      assert (Hrun : run_len x f >= 1) by (subst f; cbn [run_len]; rewrite Nat.eqb_refl; lia).
      pose proof (run_le x f) as Hrl.
      destruct (Nat.leb_spec 4 l) as [H4|H4].
      * set (f2 := skipn l f).
        assert (Hlf : length f = l + length f2) by (unfold f2; rewrite length_skipn; unfold l in *; lia).
        destruct (IH f2 ltac:(lia) ltac:(unfold f2; apply Forall_drop; exact Hb)) as (B & _ & _).
        rewrite psteps_cons. cbn [tsteps enc length]. replace (256 - (256 - l)) with l by (unfold l in *; lia).
        split; [lia|]. split; [intros _; lia|].
        intro Hn. exfalso. specialize (Hn 0). cbn [skipn] in Hn. subst f. rewrite four_run in Hn.
        apply Nat.leb_gt in Hn. unfold l in *. lia.
      * assert (Hf4 : four_at f = false).
        { subst f. rewrite four_run. apply Nat.leb_gt. unfold l in *. lia. }
        set (k := lit_scan 127 f).
        destruct (lit_scan_le 127 f) as [Hk1 Hk2]. fold k in Hk1, Hk2.
        assert (Hk0 : 1 <= k) by (apply lit_scan_pos; [lia | subst f; discriminate | exact Hf4]).
        set (f2 := skipn k f).
        assert (Hlf : length f = k + length f2) by (unfold f2; rewrite length_skipn; lia).
        destruct (IH f2 ltac:(lia) ltac:(unfold f2; apply Forall_drop; exact Hb)) as (B & B4 & E).
        rewrite psteps_cons. cbn [tsteps enc length]. rewrite length_map, length_firstn, Nat.min_l by lia.
        assert (Hd2 : length f2 = 0 -> psteps (desc fu f2) = 16).
        { intro H0. assert (Z : f2 = []) by (apply length_zero_iff_nil; exact H0).
          rewrite Z. destruct fu; reflexivity. }
        split; [|split; [intro H; rewrite Hf4 in H; discriminate|]].
        -- destruct (lit_scan_stop 127 f) as [S1|[S1|S1]]; fold k in S1.
           ++ lia.
           ++ rewrite Hd2 by lia. lia.
           ++ specialize (B4 S1). lia.
        -- intro Hn. assert (Hn2 : nofour f2) by (apply nofour_skipn; exact Hn).
           specialize (E Hn2). destruct (lit_scan_stop 127 f) as [S1|[S1|S1]]; fold k in S1.
           ++ rewrite E, Hlf, S1, ceil_add. lia.
           ++ rewrite Hd2 by lia. rewrite (ceil_small (length f)) by lia. lia.
           ++ exfalso. specialize (Hn k). rewrite S1 in Hn. discriminate.
Qed.

(* every recipe of f takes at least 4 |f| steps *)
Lemma psteps_lower : forall ks, Forall tok_ok ks -> 4 * length (delta ks) <= psteps ks.
Proof.
  induction ks as [|k ks IH]; intros Hok; [cbn; lia|]. apply Forall_cons_iff in Hok as [Hk Hok].
  rewrite psteps_cons, delta_cons, length_app. specialize (IH Hok).
  destruct k as [bs|t b]; cbn [tsteps enc dec length] in *; [rewrite length_map | rewrite repeat_length]; lia.
Qed.

Theorem describe_steps : forall f, Forall (fun x => x < 256) f ->
  2 * psteps (describe f) <= 25 * length f + 58
  /\ (nofour f -> psteps (describe f) = 11 * length f + 13 * ((length f + 126) / 127) + 16).
Proof.
  intros f Hb. destruct (desc_steps (length f) f (le_n _) Hb) as (B & _ & E). exact (conj B E).
Qed.

(* ---- Lemma 6.5 --------------------------------------------------------------------------------------------- *)
Corollary recipes_complete : forall f, Forall (fun x => x < 256) f ->
  exists ks, Forall tok_ok ks /\ delta ks = f /\ length (recipe ks) <= bound (length f).
Proof.
  intros f Hb. destruct (describe_spec f Hb) as (D & T & B & _). exists (describe f). exact (conj T (conj D B)).
Qed.

(* ---- Proposition 6.6 --------------------------------------------------------------------------------------- *)
Definition notwo (f : list nat) : Prop := forall p, S p < length f -> nth p f 0 <> nth (S p) f 0.

Lemma notwo_app : forall a b, notwo (a ++ b) -> notwo a /\ notwo b.
Proof.
  intros a b H. split.
  - intros p Hp. specialize (H p ltac:(rewrite length_app; lia)). rewrite !app_nth1 in H by lia. exact H.
  - intros p Hp. specialize (H (length a + p) ltac:(rewrite length_app; lia)).
    rewrite !app_nth2 in H by lia. replace (length a + p - length a) with p in H by lia.
    replace (S (length a + p) - length a) with (S p) in H by lia. exact H.
Qed.

Lemma token_cost : forall k, tok_ok k -> notwo (dec k) ->
  length (enc k) = length (dec k) + 1 /\ 1 <= length (dec k) <= 127.
Proof.
  intros [bs|t b] Hk Hn; cbn [tok_ok] in Hk; destruct Hk as [Hk1 Hk2]; cbn [dec] in Hn; cbn [enc dec length].
  - rewrite length_map. lia.
  - rewrite repeat_length. destruct (Nat.eq_dec (256 - t) 1) as [E|E]; [lia|].
    exfalso. apply (Hn 0); rewrite ?repeat_length; [lia|]. rewrite !nth_repeat_lt by lia. reflexivity.
Qed.

Theorem recipe_optimal : forall ks, Forall tok_ok ks -> notwo (delta ks) ->
  bound (length (delta ks)) <= length (recipe ks).
Proof.
  assert (G : forall ks, Forall tok_ok ks -> notwo (delta ks) ->
            length (recipe ks) = length (delta ks) + length ks + 1 /\ length (delta ks) <= 127 * length ks).
  { induction ks as [|k ks IH]; intros Hok Hn; [cbn; lia|].
    apply Forall_cons_iff in Hok as [Hk Hok]. rewrite delta_cons in Hn. apply notwo_app in Hn as [Hn1 Hn2].
    destruct (IH Hok Hn2) as [E1 E2]. destruct (token_cost k Hk Hn1) as [C1 C2].
    rewrite length_recipe_cons, delta_cons, length_app, C1, E1. cbn [length]. lia. }
  intros ks Hok Hn. destruct (G ks Hok Hn) as [E1 E2]. rewrite E1. unfold bound.
  assert ((length (delta ks) + 126) / 127 < S (length ks)).
  { apply Nat.Div0.div_lt_upper_bound. lia. }
  lia.
Qed.