(** ** Coverage and closure (proof.md 4.2 Lemma S, 4.3 Lemma C) *)

Local Open Scope nat_scope.

Lemma St_std N m: St (false, N, m) = K N 0 0 m.
Proof. reflexivity. Qed.

Lemma St_pk N m: St (true, N, m) = K (pkE N) (pkO N) 0 m.
Proof. reflexivity. Qed.

Lemma even_mod x: Nat.even x = true -> x mod 2 = 0.
Proof. intros H. apply Nat.even_spec in H as [q ->]. lia. Qed.

Lemma odd_mod x: Nat.even x = false -> x mod 2 = 1.
Proof.
  intros H. destruct (Nat.Even_or_Odd x) as [He|[q ->]].
  - apply Nat.even_spec in He. congruence.
  - lia.
Qed.

Ltac par x := let E := fresh "E" in
  destruct (Nat.even x) eqn:E; [apply even_mod in E | apply odd_mod in E].

(** I in linear-arithmetic form (parities as [mod 2]) *)
Lemma I_std N m: I (false, N, m) <->
  1 <= m <= N /\
  ((N+m) mod 2 = 0 -> N mod 2 = 0 -> 2*N <= 3*m+6) /\
  ((N+m) mod 2 = 0 -> N mod 2 = 1 -> N <= 2*m+3) /\
  ((N+m) mod 2 = 1 -> N <= m+3).
Proof.
  unfold I. cbv beta iota zeta. par (N+m); par N; cbv beta iota.
  all: split; [intros [H1 [_ H3]]; lia | intros H; split; [lia | split; [intros Hf; discriminate | lia]]].
Qed.

Lemma I_pk N m: I (true, N, m) <->
  1 <= m <= N /\ 2 <= N /\
  ((N+m) mod 2 = 0 -> N mod 2 = 0 -> N <= 2*m+3) /\
  ((N+m) mod 2 = 0 -> N mod 2 = 1 -> 2*N <= 5*m+6) /\
  ((N+m) mod 2 = 1 -> N <= m+3).
Proof.
  unfold I. cbv beta iota zeta. par (N+m); par N; cbv beta iota.
  all: split; [intros [H1 [H2 H3]]; specialize (H2 eq_refl); lia | intros H; split; [lia | split; [intros _; lia | lia]]].
Qed.

(** results N' = D + n, m' = n *)
Lemma close_diag t n: 1 <= n -> (t = true -> 2 <= n) -> I (t, n, n).
Proof.
  intros H1 H2. destruct t; [apply I_pk; specialize (H2 eq_refl) | apply I_std]; lia.
Qed.

Lemma close_std D n: 1 <= n ->
  (D <= 3 /\ D mod 2 = 1) \/ (D mod 2 = 0 /\ (n mod 2 = 1 -> D <= n+3) /\ (n mod 2 = 0 -> 2*D <= n+6)) ->
  I (false, D+n, n).
Proof. intros. apply I_std. lia. Qed.

Lemma close_pk D n: 1 <= n -> 2 <= D+n ->
  (D <= 3 /\ D mod 2 = 1) \/ (D mod 2 = 0 /\ (n mod 2 = 0 -> D <= n+3) /\ (n mod 2 = 1 -> 2*D <= 3*n+6)) ->
  I (true, D+n, n).
Proof. intros. apply I_pk. lia. Qed.

(** N+1 = 2^i (2k+1): the three shapes of N *)
Lemma Nfacts N k i: S N = 2^i*(2*k+1) ->
  (i = 0 /\ N = 2*k) \/ (i = 1 /\ N = 4*k+1) \/ (2 <= i /\ (S N) mod 4 = 0 /\ 8*k+4 <= S N).
Proof.
  intros H. destruct i as [|[|i']].
  - left. rewrite Nat.pow_0_r in H. lia.
  - right; left. rewrite Nat.pow_1_r in H. lia.
  - right; right. rewrite !Nat.pow_succ_r' in H. pose proof (pow_pos i').
    assert (H1: S N = 4 * (2^i'*(2*k+1))) by lia.
    assert (H2: 2*k+1 <= 2^i'*(2*k+1)) by nia.
    remember (2^i'*(2*k+1)) as Y. clear HeqY H. lia.
Qed.

Ltac nfacts N k i Hk :=
  destruct (Nfacts N k i Hk) as [[-> ->]|[[-> ->]|[?Hi [?H4 ?H8]]]]; clear Hk.

Lemma log2_facts k: 1 <= k -> exists l r, k = 2^l + r /\ r < 2^l.
Proof.
  intros Hk. exists (Nat.log2 k), (k - 2^(Nat.log2 k)).
  pose proof (Nat.log2_spec k ltac:(lia)) as [H1 H2].
  rewrite Nat.pow_succ_r' in H2. lia.
Qed.

Lemma pow_ge2 u: 1 <= u -> 2 <= 2^u.
Proof. intros Hu. destruct u as [|u]; [lia|]. rewrite Nat.pow_succ_r'. pose proof (pow_pos u). lia. Qed.

Lemma pow_ge4 u: 2 <= u -> 4 <= 2^u.
Proof.
  intros Hu. replace u with (S (S (u-2))) by lia. rewrite !Nat.pow_succ_r'.
  pose proof (pow_pos (u-2)). lia.
Qed.

Lemma pow_par h: (h = 0 /\ 2^h = 1) \/ 2^h mod 2 = 0.
Proof. destruct h as [|h]; [left; split; reflexivity | right; rewrite Nat.pow_succ_r'; lia]. Qed.

(** the gen-2 exponents of S3/P4: v = 2^u (2g+1) with 1 <= v < 2^l and v even *)
Lemma split_v l v: 1 <= v -> v < 2^l -> v mod 2 = 0 ->
  exists h u g, l = S (h+u) /\ v = 2^u*(2*g+1) /\ 1 <= u /\ g < 2^h.
Proof.
  intros Hv1 Hv2 Hv3.
  destruct (lowbit v) as [u [g Hv]]; [lia|].
  pose proof (pow_odd_pos u g) as Pug.
  assert (Hul: u < l). { apply (Nat.pow_lt_mono_r_iff 2); lia. }
  exists (l - S u), u, g.
  assert (El: l = S (l - S u + u)) by lia.
  split; [exact El|]. split; [exact Hv|].
  pose proof (pow_pos u) as Pu.
  split.
  - destruct u as [|u']; [|lia]. rewrite Nat.pow_0_r in Hv. lia.
  - set (h := l - S u) in *. clearbody h. subst l.
    rewrite Nat.pow_succ_r', Nat.pow_add_r in Hv2.
    assert (H1: 2^u*(2*g+1) < 2^u*(2*2^h)) by lia.
    apply Nat.mul_lt_mono_pos_l in H1; lia.
Qed.

(** closure of S1, P1, P2 (proof.md 4.3) *)
Lemma C_S1 N m i k: I (false, N, m) -> S N = 2^i*(2*k+1) -> 1 <= k ->
  I (true, 2*k+(1+i+m), 1+i+m).
Proof.
  intros HI Hk Hk1. apply I_std in HI. apply I_pk.
  nfacts N k i Hk; lia.
Qed.

Lemma C_P1 N m i l: I (true, N, m) -> S N = 2^i*(2*(2^l+0)+1) ->
  I (false, 2^l+(1+i+m), 1+i+m).
Proof.
  intros HI Hk. apply I_pk in HI. apply I_std.
  pose proof (pow_pos l). pose proof (pow_par l) as Hl.
  rewrite Nat.add_0_r in Hk. set (Z := 2^l) in *. clearbody Z.
  nfacts N Z i Hk; lia.
Qed.

Lemma C_P2 N m i l r: I (true, N, m) -> S N = 2^i*(2*(2^l+r)+1) -> r < 2^l ->
  I (false, 2*r+(1+i+m), 1+i+m).
Proof.
  intros HI Hk Hr. apply I_pk in HI. apply I_std.
  set (Z := 2^l) in *. clearbody Z.
  nfacts N (Z+r) i Hk; lia.
Qed.

(** coverage of S3 and P4 (proof.md 4.2): the dirty cases are i = 0 (S3), i <= 1 (P4), class P,
    in the final gap, with T even *)
Lemma Cov_S3 N m i l r: I (false, N, m) -> S N = 2^i*(2*(2^l+r)+1) -> 1 <= r -> r < 2^l ->
  2*(2^l+r)+(1+i+m)+1 < 4*2^l ->
  i = 0 /\ m mod 2 = 0 /\ 2*r < 2^l /\ 3*2^l < 2*(2^l+r)+(1+i+m)+1 /\ 4*(2^l+r) <= 3*m+6.
Proof.
  intros HI Hk Hr1 Hr Hlt. apply I_std in HI.
  set (Z := 2^l) in *. clearbody Z.
  nfacts N (Z+r) i Hk; lia.
Qed.

Lemma Cov_P4 N m i l r: I (true, N, m) -> S N = 2^i*(2*(2^l+r)+1) -> 1 <= r -> r < 2^l ->
  2*r+(1+i+m)+1 < 2*2^l ->
  i <= 1 /\ (2*r+(1+i+m)+1) mod 2 = 0 /\ 2*r < 2^l /\ 2^l < 2*r+(1+i+m)+1 /\ 2^l < 1+i+m.
Proof.
  intros HI Hk Hr1 Hr Hlt. apply I_pk in HI.
  set (Z := 2^l) in *. clearbody Z.
  nfacts N (Z+r) i Hk; lia.
Qed.

Lemma step_std N m: I (false, N, m) -> exists s', St (false, N, m) -->+ St s' /\ I s'.
Proof.
  intros HI. rewrite St_std.
  assert (Hm: 1 <= m <= N) by (apply I_std in HI; tauto).
  destruct (lowbit (S N)) as [i [k Hk]]; [lia|].
  destruct (Nat.eq_dec k 0) as [->|Hk0].
  { (* S0 *)
    exists (false, 1+i+m, 1+i+m). rewrite St_std. split.
    - replace N with (2^i-1) by lia. apply BS_S0.
    - apply close_diag; [lia | intros Hf; discriminate Hf]. }
  destruct (log2_facts k ltac:(lia)) as [l [r [-> Hr]]].
  replace N with ((2*(2^l+r)+1)*2^i-1) by lia.
  pose proof (pow_pos l) as Pl.
  assert (E2: 2^(S l) = 2*2^l) by apply Nat.pow_succ_r'.
  assert (E4: 2^(S (S l)) = 4*2^l) by (rewrite !Nat.pow_succ_r'; lia).
  destruct (Nat.eq_dec r 0) as [->|Hr0].
  { (* S1, k = 2^l *)
    exists (true, 2*(2^l+0)+(1+i+m), 1+i+m). rewrite St_pk. split.
    - apply BS_S1a with (l := l). lia.
    - apply (C_S1 N); [exact HI | exact Hk | lia]. }
  destruct (lt_eq_lt_dec (2*(2^l+r)+(1+i+m)+1) (2^(S (S l)))) as [[Hlt|Heq]|Hgt].
  3: { (* S1, T > 2^(l+2) *)
    exists (true, 2*(2^l+r)+(1+i+m), 1+i+m). rewrite St_pk. split.
    - clear Hk. apply BS_S1b with (l := l); lia.
    - apply (C_S1 N); [exact HI | exact Hk | lia]. }
  2: { (* S2 *)
    exists (false, l+4+(1+i+m), l+4+(1+i+m)). rewrite St_std. split.
    - clear Hk. apply BS_S2; lia.
    - apply close_diag; [lia | intros Hf; discriminate Hf]. }
  (* S3 *)
  rewrite E4 in Hlt.
  destruct (Cov_S3 N m i l r HI Hk ltac:(lia) Hr Hlt) as [-> [Hm2 [Hr2 [HT3 H43]]]].
  clear HI Hk Hm.
  assert (HZ: 2^l mod 2 = 0).
  { destruct l as [|l']; [rewrite Nat.pow_0_r in *; lia | rewrite Nat.pow_succ_r'; lia]. }
  destruct (split_v l (2*(2^l+r)+(1+0+m)+1 - 3*2^l)) as [h [u [g [-> [Hv [Hu1 Hg]]]]]];
    [lia|lia|lia|].
  assert (E1: 2^(S (h+u)) = 2*(2^h*2^u)) by (rewrite Nat.pow_succ_r', Nat.pow_add_r; reflexivity).
  assert (E3: 2^(h+u) = 2^h*2^u) by apply Nat.pow_add_r.
  assert (E5: 2^(S h) = 2*2^h) by apply Nat.pow_succ_r'.
  pose proof (pow_pos h) as Ph. pose proof (pow_pos u) as Pu.
  assert (HU: 2*2^h <= 2^h*2^u) by (pose proof (pow_ge2 u Hu1); nia).
  assert (HU2: 2 <= u -> 4*2^h <= 2^h*2^u) by (intros Hu2; pose proof (pow_ge4 u Hu2); nia).
  pose proof (pow_par h) as Hh.
  (* n = 1+u+(1+0+m) > 2^l >= 4*2^h, n odd if u = 1 *)
  assert (Hn: 2^(S (h+u)) < 1+u+(1+0+m)) by lia.
  assert (Hv': 2*r+(1+0+m)+1 = 2^(S (h+u)) + 2^u*(2*g+1)) by lia.
  clear Hv HZ E4 E2 Pl HT3 H43 Hlt.
  destruct (Nat.eq_dec g 0) as [->|Hg0].
  - exists (false, 3*2^h+(1+u+(1+0+m)), 1+u+(1+0+m)). rewrite St_std. split.
    + apply BS_S3a; lia.
    + apply close_std; [lia|].
      destruct Hh as [[-> Hh1]|Hh]; [left; lia | right; lia].
  - exists (false, 2^(S h)+2*g+(1+u+(1+0+m)), 1+u+(1+0+m)). rewrite St_std. split.
    + apply BS_S3b; lia.
    + apply close_std; [lia|]. right. lia.
Qed.

Lemma step_pk N m: I (true, N, m) -> exists s', St (true, N, m) -->+ St s' /\ I s'.
Proof.
  intros HI. rewrite St_pk.
  assert (Hm: 1 <= m <= N /\ 2 <= N) by (apply I_pk in HI; tauto).
  destruct (lowbit (S N)) as [i [k Hk]]; [lia|].
  destruct (Nat.eq_dec k 0) as [->|Hk0].
  { (* P0 *)
    assert (Hi: 2 <= i).
    { destruct (Nfacts N 0 i Hk) as [[_ HN]|[[_ HN]|[Hi _]]]; lia. }
    exists (false, 2+i+m, 2+i+m). rewrite St_std. split.
    - replace N with (2^i-1) by lia. apply BS_P0, Hi.
    - apply close_diag; [lia | intros Hf; discriminate Hf]. }
  destruct (log2_facts k ltac:(lia)) as [l [r [-> Hr]]].
  replace N with ((2*(2^l+r)+1)*2^i-1) by lia.
  pose proof (pow_pos l) as Pl.
  assert (E2: 2^(S l) = 2*2^l) by apply Nat.pow_succ_r'.
  destruct (Nat.eq_dec r 0) as [->|Hr0].
  { (* P1 *)
    exists (false, 2^l+(1+i+m), 1+i+m). rewrite St_std. split.
    - apply BS_P1.
    - apply (C_P1 N); [exact HI | exact Hk]. }
  destruct (lt_eq_lt_dec (2*r+(1+i+m)+1) (2^(S l))) as [[Hlt|Heq]|Hgt].
  3: { (* P2 *)
    exists (false, 2*r+(1+i+m), 1+i+m). rewrite St_std. split.
    - clear Hk. apply BS_P2; lia.
    - apply (C_P2 N m i l); [exact HI | exact Hk | exact Hr]. }
  2: { (* P3 *)
    exists (false, l+2+(1+i+m), l+2+(1+i+m)). rewrite St_std. split.
    - clear Hk. apply BS_P3; lia.
    - apply close_diag; [lia | intros Hf; discriminate Hf]. }
  (* P4 *)
  rewrite E2 in Hlt.
  destruct (Cov_P4 N m i l r HI Hk ltac:(lia) Hr Hlt) as [Hi1 [HT2 [Hr2 [HT1 Hb]]]].
  clear HI Hk Hm.
  assert (HZ: 2^l mod 2 = 0).
  { destruct l as [|l']; [rewrite Nat.pow_0_r in *; lia | rewrite Nat.pow_succ_r'; lia]. }
  destruct (split_v l (2*r+(1+i+m)+1 - 2^l)) as [h [u [g [-> [Hv [Hu1 Hg]]]]]];
    [lia|lia|lia|].
  assert (E1: 2^(S (h+u)) = 2*(2^h*2^u)) by (rewrite Nat.pow_succ_r', Nat.pow_add_r; reflexivity).
  assert (E3: 2^(h+u) = 2^h*2^u) by apply Nat.pow_add_r.
  assert (E5: 2^(S h) = 2*2^h) by apply Nat.pow_succ_r'.
  pose proof (pow_pos h) as Ph. pose proof (pow_pos u) as Pu.
  assert (HU: 2*2^h <= 2^h*2^u) by (pose proof (pow_ge2 u Hu1); nia).
  pose proof (pow_par h) as Hh.
  assert (Hv': 2*r+(1+i+m)+1 = 2^(S (h+u)) + 2^u*(2*g+1)) by lia.
  clear Hv HZ E2 Pl HT1 HT2 Hlt.
  destruct (Nat.eq_dec g 0) as [->|Hg0].
  - exists (true, 3*2^h+(1+u+(1+i+m)), 1+u+(1+i+m)). rewrite St_pk. split.
    + apply BS_P4a; lia.
    + apply close_pk; [lia|lia|].
      destruct Hh as [[-> Hh1]|Hh]; [left; lia | right; lia].
  - exists (true, 2^(S h)+2*g+(1+u+(1+i+m)), 1+u+(1+i+m)). rewrite St_pk. split.
    + apply BS_P4b; lia.
    + apply close_pk; [lia|lia|]. right. lia.
Qed.

Lemma step_I s: I s -> exists s', St s -->+ St s' /\ I s'.
Proof.
  destruct s as [[[|] N] m].
  - apply step_pk.
  - apply step_std.
Qed.