(* Every record image has a description: translate inverts describe. *)
From Stdlib Require Import List Bool Arith Lia.
Import ListNotations.
From TH Require Import TwoHead Decoder.

Definition seg (r : list bool) : list bool :=
  match find_codon candidates (firstn 8 r) with
  | Some c => c ++ skipn 8 r
  | None => ESC ++ r
  end.

Lemma describe_eq : forall X, describe X = concat (map seg (chunks (length X) X)) ++ ENDC.
Proof. reflexivity. Qed.

Lemma find_codon_some : forall cs g c, find_codon cs g = Some c ->
  exists k a b, In (k, a, b) cs /\ gate_bits k a b = g /\ c = codon_bits k a b.
Proof.
  induction cs as [|[[k a] b] cs IH]; intros g c H; cbn [find_codon] in H; [discriminate|].
  destruct (lbeq (gate_bits k a b) g) eqn:E.
  - apply some_inj in H. subst c. exists k, a, b.
    split; [left; reflexivity|]. split; [apply lbeq_true; exact E | reflexivity].
  - destruct (IH g c H) as (k' & a' & b' & Hin & Hg & Hc).
    exists k', a', b'. split; [right; exact Hin|]. split; assumption.
Qed.

Lemma candidates_kind : forall k a b, In (k, a, b) candidates -> k <= 2.
Proof.
  intros k a b H. unfold candidates in H.
  repeat (destruct H as [H|H]; [injection H as <- <- <-; lia|]). destruct H.
Qed.

Lemma translate_codon : forall f k a b X, k <= 2 ->
  translate_aux (S f) (codon_bits k a b ++ X)
  = option_map (fun t => gate_bits k a b ++ firstn 18 X ++ t) (translate_aux f (skipn 18 X)).
Proof.
  intros f k a b X Hk. unfold codon_bits. cbn [app]. rewrite translate_aux_S.
  destruct k as [|[|[|k]]]; [| | | lia]; cbn [Nat.eqb orb andb]; reflexivity.
Qed.

Lemma translate_esc : forall f r rest, length r = 26 ->
  translate_aux (S f) (ESC ++ r ++ rest) = option_map (fun t => r ++ t) (translate_aux f rest).
Proof.
  intros f r rest Hr. unfold ESC. cbn [app]. rewrite translate_aux_S. cbn [andb].
  replace (Nat.eqb (length (r ++ rest)) 0) with false
    by (symmetry; apply Nat.eqb_neq; rewrite length_app; lia).
  rewrite (firstn_app_exact r rest 26 Hr), (skipn_app_exact r rest 26 Hr). reflexivity.
Qed.

Lemma seg_translate : forall f r rest, length r = 26 ->
  translate_aux (S f) (seg r ++ rest) = option_map (fun t => r ++ t) (translate_aux f rest).
Proof.
  intros f r rest Hr. unfold seg.
  destruct (find_codon candidates (firstn 8 r)) as [c|] eqn:Ef.
  - destruct (find_codon_some _ _ _ Ef) as (k & a & b & Hin & Hg & ->).
    assert (H18 : length (skipn 8 r) = 18) by (rewrite length_skipn; lia).
    rewrite <- app_assoc, translate_codon by exact (candidates_kind k a b Hin).
    rewrite (firstn_app_exact _ rest 18 H18), (skipn_app_exact _ rest 18 H18), Hg.
    destruct (translate_aux f rest) as [t|]; cbn [option_map]; [|reflexivity].
    rewrite app_assoc, firstn_skipn. reflexivity.
  - rewrite <- app_assoc. apply translate_esc. exact Hr.
Qed.

Lemma seg_long : forall r, 1 <= length (seg r).
Proof.
  intro r. unfold seg. destruct (find_codon candidates (firstn 8 r)) as [c|] eqn:Ef.
  - destruct (find_codon_some _ _ _ Ef) as (k & a & b & _ & _ & ->).
    rewrite length_app. unfold codon_bits. cbn [length]. lia.
  - unfold ESC. rewrite length_app. cbn [length]. lia.
Qed.

Lemma chunks_S : forall f l, l <> [] -> chunks (S f) l = firstn 26 l :: chunks f (skipn 26 l).
Proof. intros f [|x l] H; [congruence|reflexivity]. Qed.

Lemma chunks_nil : forall f, chunks f [] = [].
Proof. intros [|f]; reflexivity. Qed.

Lemma describe_aux : forall n X f f', length X = 26 * n -> n <= f ->
  length (concat (map seg (chunks f X)) ++ ENDC) <= f' ->
  translate_aux f' (concat (map seg (chunks f X)) ++ ENDC) = Some X.
Proof.
  induction n as [|n IH]; intros X f f' HX Hf Hf'.
  - destruct X; [|discriminate]. rewrite chunks_nil in *. cbn [map concat app] in *.
    destruct f' as [|f']; [cbn [length ENDC] in Hf'; lia|].
    unfold ENDC. rewrite translate_aux_S. reflexivity.
  - destruct f as [|f]; [lia|].
    destruct (split_at X 26 ltac:(lia)) as (r & X' & -> & Hr).
    rewrite length_app, Hr in HX.
    assert (Hne : r ++ X' <> []) by (destruct r; [discriminate | intro H; discriminate]).
    rewrite chunks_S in Hf' |- * by exact Hne.
    rewrite (firstn_app_exact r X' 26 Hr), (skipn_app_exact r X' 26 Hr) in Hf' |- *.
    cbn [map concat] in Hf' |- *. rewrite <- app_assoc in Hf' |- *.
    rewrite length_app in Hf'. pose proof (seg_long r).
    destruct f' as [|f']; [lia|].
    rewrite seg_translate by exact Hr.
    rewrite (IH X' f f') by lia. reflexivity.
Qed.

(* Every image of n records has the description describe X, which translate reads back. *)
Theorem describe_complete : forall X n, length X = 26 * n -> translate (describe X) = Some X.
Proof.
  intros X n HX. unfold translate. rewrite describe_eq.
  apply describe_aux with (n := n); [exact HX | lia | lia].
Qed.