Set Implicit Arguments.
Require Import Coq.Init.Nat.

Require Import BasicDefs.
Require Import FinType.
Require Import Seqs.

Open Scope string_scope.

Structure NFA (X:finType) := mknfa
   { state: finType;
     transition : state -> X -> state -> Prop;
     transition_dec : forall s1 a s2, dec(transition s1 a s2);
     initial_state : state -> Prop;
     initial_state_dec : forall s, dec(initial_state s);
     final_state: state -> Prop;
     final_state_dec: forall s, dec(final_state s)
   }.

Arguments mknfa [X] [state] transition {transition_dec} initial_state {initial_state_dec} final_state {final_state_dec}.

Arguments transition [X] [aut] _ _ _ : rename.
Arguments initial_state [X] [aut] _ : rename.
Arguments final_state [X] [aut] _ : rename.
Arguments transition_dec [X] [aut] _ _ _ : rename.
Arguments initial_state_dec [X] [aut] _ : rename.
Arguments final_state_dec [X] [aut] _ : rename.

Hint Resolve final_state_dec.
Hint Resolve initial_state_dec.
Hint Resolve transition_dec.

(* Dont ask me why Coq sometimes needs this  *)
Instance transition_dec_public (X:finType)(aut:NFA X): forall s x s', dec (transition (aut:=aut) s x s').
Proof. auto. Defined.
Instance initial_dec_public (X:finType)(aut:NFA X): forall s, dec (initial_state (aut:=aut) s).
Proof. auto. Defined.
Instance final_dec_public (X:finType)(aut:NFA X): forall s, dec (final_state (aut:=aut) s).
Proof. auto. Defined.

Definition univ_aut {X:finType} := mknfa (fun (s:finType_unit) (x:X) (s':finType_unit) => True) (fun _ => True) (fun _ => True).

Section TransitionGraph.
  Context {X : finType}.
  Context {aut: NFA X}.

  Definition Run := Seq (state aut).


  Implicit Types (r : Run ) (w: Seq X) (v: String X) (s: state aut) (x:X).

  Definition valid_run r w:= forall n, transition (r n) (w n) (r (S n)).
  Definition valid_path r v:= forall n, n <= |v| -> transition (r n) (v [n]) (r (S n)).

  Lemma valid_run_extensional r r' w w': valid_run r w -> r === r' -> w === w' -> valid_run r' w'.
  Proof.
    intros V R W n. now rewrite <- W, <- 2R.
  Qed.

  Lemma valid_path_extensional r r' v v': valid_path r v -> r === r' -> v == v' -> valid_path r' v'.
  Proof.
    intros V E [LE WE] n L. rewrite <-LE in L. rewrite <-WE by auto. rewrite <-!E. now apply V.
  Qed.

  Global Instance valid_run_proper : Proper (@seq_equal (state aut) ==> @seq_equal X ==> iff) valid_run.
  Proof.
    intros r r' EqR w w' EqW. split; intros V; now eapply valid_run_extensional.
  Qed.

  (* TODO seq_equal is over approx *)
  Global Instance valid_path_proper : Proper (@seq_equal (state aut) ==> strings_equal ==> iff) valid_path .
  Proof.
    intros r r' Er w w' Ew. split; intros V.
    - apply (valid_path_extensional V Er Ew).
    - symmetry in Er, Ew. apply (valid_path_extensional V Er Ew).
  Qed.

  Lemma valid_run_is_path_from_begin r w: valid_run r w -> forall n, valid_path r (mkstring w n).
  Proof. now intros. Qed.


  Lemma valid_path_decrease' r v n: valid_path r (mkstring v n) -> forall k, k <= n -> valid_path r (mkstring v k).
  Proof.

    intros V k' L m Lm. cbn. apply V. cbn in *. omega.
  Qed.

   Lemma valid_path_decrease r v : valid_path r v -> forall k, k <= |v| -> valid_path r (drop_string_end k v).
  Proof.
    destruct v. unfold drop_string_end. cbn. intros V k L.
    apply (valid_path_decrease' V). comega.
  Qed.

  Lemma valid_run_drop r w: forall n, valid_run r w -> valid_run (drop n r) (drop n w).
  Proof.
    intros n V m. now rewrite !drop_correct, <-plus_n_Sm.
  Qed.

  Lemma valid_path_drop r v k: valid_path r v -> k <= |v| -> valid_path (drop k r) (drop_string_begin k v).
  Proof.
    intros V L n Lk. unfold drop_string_begin. cbn in *.
    rewrite !drop_correct, <-plus_n_Sm. apply V. omega.
  Qed.

  Lemma valid_run_is_path_everywhere r w: valid_run r w -> forall n m, valid_path (drop n r) (extract n m w).
  Proof.
    intros V n m k L. unfold extract. cbn.
    now apply valid_run_drop.
  Qed.

  Lemma valid_path_is_path_everywhere r v: valid_path r v -> forall n m, n < m <= S(|v|) -> valid_path (drop n r) (extract n m v).
  Proof.
    intros V n m L k L'. unfold extract. cbn in *.
    apply valid_path_drop; oauto.
  Qed.

  Lemma prepend_step_valid_path s x r v : valid_path r v -> transition s x (r 0) -> valid_path (prepend s r) (mkstring (prepend x v) (S(|v|))).
  Proof.
    intros V T n L. destruct n.
    - now rewrite !prepend_first, prepend_rest.
    - rewrite !prepend_rest. apply V. cbn in L. omega.
  Qed.

  Lemma valid_concat_paths r v r' v': valid_path r v -> r' 0 = r (S (|v|)) -> valid_path r' v' ->
                      valid_path (mkstring r (|v|) +++ r') (v ++ v').
  Proof.
   destruct v as [v k]. revert v r v' r'.
  induction k; intros w r r' w' P C P'. cbn in C.
  - apply prepend_step_valid_path; auto. rewrite C. cbn. now apply P.
  - cbn in *. rewrite concat_strings_step. apply IHk.
    + apply (valid_path_decrease' P). comega.
    + reflexivity.
    + apply prepend_step_valid_path; oauto. rewrite C. now apply P.
  Qed.

  Lemma valid_prepend_path' (r':String (state aut)) v r w : |r'| = |v| -> (forall n, n < |v| -> transition (r'[n]) (v[n]) (r' [S n]))
                                  -> transition (r'[|v|]) (v[|v|]) (r 0)-> valid_run r w->
                               valid_run ( r' +++ r) (v +++ w).
  Proof.
    intros E V C R. intros n.
    decide (n < |v|).
    - rewrite 3prepend_string_begin by omega. now apply V.
    - decide (n = |v|).
      + subst n. rewrite 2prepend_string_begin by omega. rewrite_eq (S(|v|)=S(|v|) + 0). rewrite <-E, prepend_string_rest', E. apply C.
      + rewrite 3prepend_string_rest by omega. rewrite_eq (S n - S(|r'|) = S (n - S(|r'|))). rewrite E. apply R.
  Qed.

  Lemma valid_prepend_path r' v r w : valid_path r' v -> (r 0) = (r' (S (|v|))) -> (valid_run r w) ->
                               valid_run ((mkstring r' (|v|) ) +++ r) (v +++ w).
  Proof.
    intros V C R. intros n.
    apply (valid_concat_paths (r' := r) (v := v) (r := r' ) (v' := mkstring w n)); auto.
    - now apply valid_run_is_path_from_begin.
    - comega.
  Qed.

  Lemma valid_path_cut r v i j: valid_path r v -> S i < j <= |v| -> r (S i) = r j -> valid_path (cut i j r) (cut_string i j v).
  Proof.
    intros V L E. unfold cut, cut_string.
    apply valid_concat_paths.
    - apply (valid_path_decrease' V). omega.
    - rewrite drop_correct. cbn. rewrite E. now f_equal.
    - apply valid_path_drop; oauto.
  Qed.

  Section InfConcat.

  Variable strings: Seq (String X).
  Variable runs: Seq (String (state aut)).

  Lemma concat_inf_is_valid'':
        (forall n, |strings n| = |runs n| /\ (forall k, k < |strings n| -> transition ((runs n)[k])((strings n)[k])((runs n)[S k])) /\ transition ((runs n)[|runs n|]) ( (strings n) [|runs n|]) ((runs (S n)) [0]))
        -> valid_run (concat_inf runs) (concat_inf strings).
  Proof.
    intros V n. unfold concat_inf.
    rewrite (concat_inf_index_equal (f:= strings) (g:= runs)) by apply V.
    destruct (concat_inf_indices runs n) as [wi ni] eqn:H.
    destruct (concat_inf_index_step_correct H) as [H' | [H'1 H'2]].
    - rewrite H'. simpl. apply V.
      rewrite <-(concat_inf_index_equal (f:= strings) (g:= runs)) in H' by apply V.
      pose (concat_inf_index_in_bounds H'). omega.
    - rewrite H'1, H'2. simpl.
      destruct (V wi) as [L [V' B]].
      apply B.
  Qed.

  Lemma concat_inf_is_valid:
        (forall n, |strings n| = |runs n| /\ valid_path (seq(runs n)) (strings n) /\ transition ((runs n)[|runs n|]) ( (strings n) [|runs n|]) ((runs (S n)) [0]))
        -> valid_run (concat_inf runs) (concat_inf strings).
  Proof.
    intros V. apply concat_inf_is_valid''; auto.
    intros n. specialize (V n). destruct V as [V1 [V2 V3]]. repeat split; auto.
    intros k L. apply V2. omega.
  Qed.

  Lemma concat_inf_is_valid':
        (forall n, |strings n| = |runs n| /\ valid_path (seq(runs n)) (strings n) /\ ((runs n)[S(|runs n|)]) = ((runs (S n)) [0]))
        -> valid_run (concat_inf runs) (concat_inf strings).
  Proof.
    intros V. apply concat_inf_is_valid.
    intros n. destruct (V n) as [L [P E]]. repeat split; auto.
    rewrite <-E. apply P. omega.
  Qed.

  End InfConcat.

  Definition path r v s s' := valid_path r v /\ r 0 = s /\ r (S (|v|)) = s'.
  Definition reachable s s' := exists r v, path r v s s'.

  Lemma path_extensional r v v' s s': path r v s s' -> v == v' -> path r v' s s'.
  Proof.
    intros [V [B E]] Ev.
    repeat split; auto; now rewrite <-Ev.
  Qed.

  Global Instance path_proper r : Proper (strings_equal ==> eq ==> eq ==> iff) (path r).
  Proof.
    intros v v' Ev s_1 s_1' E_1 s_2 s_2' E_2. subst s_1' s_2'.
    split; intros P; apply (path_extensional P); [assumption | now symmetry].
  Qed.

  Lemma valid_run_reachable r w: valid_run r w -> forall n m, n < m -> reachable (r n) (r m).
  Proof.
    intros V n m L. exists (drop n r); exists (extract n m w). split.
    - now apply valid_run_is_path_everywhere.
    - split; rewrite drop_correct; f_equal; comega.
  Qed.

  Lemma valid_path_reachable r v: valid_path r v -> forall i j, i < j <= S(|v|) -> reachable (r i) (r j).
  Proof.
   intros V i j L. exists (drop i r), (extract i j v). split.
   - now apply valid_path_is_path_everywhere.
   - split; rewrite drop_correct; f_equal; comega.
  Qed.

  Lemma concat_paths r r' v v' s s'' s': path r v s s'' -> path r' v' s'' s' -> path (mkstring r (|v|) +++ r') (v ++ v') s s'.
  Proof.
    intros [V [B E]] [V'[B' E']]. repeat split.
    - apply valid_concat_paths; auto. now rewrite B',E.
    - now rewrite prepend_string_begin by omega.
    - rewrite prepend_string_rest by comega. rewrite <-E'. f_equal. comega.
  Qed.

  Lemma split_path r u v s s' :path r (u ++ v) s s' -> exists s'', path r u s s'' /\ path (drop (S(|u|)) r) v s'' s'.
  Proof.
    intros [V [B E]]. exists (r (S(|u|))). repeat split;auto.
    - rewrite revert_concat_first with (v0 := v).
      apply valid_path_decrease; cbn; oauto.
    - rewrite (revert_concat_second u v).
      apply valid_path_drop; cbn ; oauto.
    - rewrite drop_correct. f_equal. omega.
    - rewrite drop_correct. now rewrite <-E.
  Qed.

  Lemma cycle_omega_is_valid_run r v s: path r v s s -> valid_run ((mkstring r (|v|)) to_omega) (v to_omega).
  Proof.
    intros [V [R0 RSk]].
    unfold omega_iteration.
    apply concat_inf_is_valid'. intros _. repeat split; auto.
    cbn. now rewrite R0.
  Qed.

End TransitionGraph.

Arguments Run [X] _.

Section DecConnected.

  Context {X: finType}.
  Context {aut : NFA X}.

  Implicit Types (r: Run aut) (w: Seq X) (v: String X) (s: state aut).

  Notation PathTriple := (prod (prod nat (Run aut)) (Seq X) : Type).
  Notation run T := (snd (fst T)).
  Notation len T := (fst (fst T)).
  Notation str T := (snd T).
  Notation valid_triple Z := (valid_path (run Z) (mkstring (str Z) (len Z))).

  Definition cut_loop (Z: PathTriple) : PathTriple := match Z with (k, r, w) =>
      match (can_find_duplicate k r) with
      | inleft (existT _ n1 (existT _ n2 _)) => match n1 with
              | 0 => (k - n2, drop n2 r, drop n2 w)
              | S n1' => (k - n2 + n1, cut (pred n1) (n2) r, cut (pred n1) (n2) w)
               end
      | _ => (k,r, w)
   end end.

  Fixpoint rem_loops' k Z := match k with
     | 0 => Z
     | S k => cut_loop (rem_loops' k Z)
  end.
  Definition rem_loops r v := rem_loops' (|v|) (|v|,r, seq v).

  Lemma cut_loop_valid Z: valid_triple Z -> valid_triple (cut_loop Z).
  Proof.
    destruct Z as [[k r] w]. intros V. cbn in V. simpl.
    destruct (can_find_duplicate (k) r) as [[i [j [O E]]] | _].
    - destruct i eqn:Ei.
      + cbn. apply (valid_path_drop (k:=j) V). comega.
      + rewrite_eq (pred (S n) = n). rewrite_eq (k - j + S n = S k - j + n).
        rewrite !pair_proj1, !pair_proj2, <-cut_string_cut by omega.
        now apply (valid_path_cut ).
    - assumption.
  Qed.

  Lemma cut_loop_decrease Z: len (cut_loop Z) < len Z \/ len (cut_loop Z) <= Cardinality (state aut) .
  Proof.
   destruct Z as [[k r] w]. unfold cut_loop. destruct (can_find_duplicate) as [ [[|i] [j[L _]]]|D].
   - left. comega.
   - left. comega.
   - now right.
  Qed.

  Lemma cut_loop_preserves_first_state Z: run (cut_loop Z) 0 = run Z 0.
  Proof.
    destruct Z as [[k r] w]. unfold cut_loop. destruct (can_find_duplicate k r) as [[[|i] [ j [ L E]]]|D]; simpl.
    - rewrite drop_correct. now rewrite_eq (j + 0= j).
    - now rewrite cut_front by oauto.
    - reflexivity.
  Qed.

  Lemma cut_loop_preserves_last_state Z: run (cut_loop Z) (S (len (cut_loop Z))) = run Z (S (len Z)).
  Proof.
    destruct Z as [[k r] w]. unfold cut_loop. destruct (can_find_duplicate k r) as [[[|i] [ j [ L E]]]|D]; simpl.
    - rewrite drop_correct. f_equal. omega.
    - rewrite_eq (S(k - j + S i) = S i + (S k - j)). rewrite cut_rest. f_equal. omega.
    - reflexivity.
  Qed.

  Lemma rem_loops'_valid k Z: valid_triple Z -> valid_triple (rem_loops' k Z).
  Proof.
    intros V. induction k; cbn; auto using cut_loop_valid.
  Qed.

  Lemma rem_loops_valid r v : valid_path r v -> valid_triple (rem_loops r v).
  Proof.
    intros P. unfold rem_loops. apply rem_loops'_valid. cbn. apply P.
  Qed.

  Lemma rem_loops'_length k Z : len (rem_loops' k Z) <= max ( len Z - k) (Cardinality (state aut)).
  Proof.
    induction k; cbn.
    - apply max_le_left. omega.
    - destruct (cut_loop_decrease (rem_loops' k Z)).
      + destruct (max_le IHk) as [D|D].
        * apply max_le_left. omega.
        * apply max_le_right. omega.
      + now apply max_le_right.
  Qed.

  Lemma rem_loops_length r v: len (rem_loops r v) <= Cardinality(state aut ).
  Proof.
    unfold rem_loops.
    pose (L := rem_loops'_length (| v |) (| v |, r, seq v)). cbn in *.
    rewrite_eq ((|v|)-(|v|) = 0) in L. now rewrite Nat.max_0_l in L.
  Qed.

  Lemma rem_loops'_preserves_ends k Z: run (rem_loops' k Z) 0 = run Z 0 /\ run (rem_loops' k Z) (S (len (rem_loops' k Z))) = run Z (S (len Z)).
  Proof.
    induction k.
    - cbn. tauto.
    - cbn in *. destruct IHk as [H1 H2]. split.
      + now rewrite cut_loop_preserves_first_state, H1.
      + now rewrite cut_loop_preserves_last_state, H2.
  Qed.

  Lemma rem_loops_preserves_ends r v: run (rem_loops r v) 0 = r 0 /\ run (rem_loops r v) (S (len (rem_loops r v))) = r (S (|v|)).
  Proof.
    unfold rem_loops. apply rem_loops'_preserves_ends.
  Qed.

  Lemma path_length_bounded s s': (exists r v, |v| <= (Cardinality (state aut)) /\ path (aut:=aut) r v s s') <-> reachable s s'.
  Proof.
    split.
    - firstorder.
    - intros [r [v [V [B E]]]].
      exists ( run (rem_loops r v)), (mkstring (str (rem_loops r v)) (len (rem_loops r v))). split.
      + cbn. apply rem_loops_length.
      + split.
        * now apply rem_loops_valid.
        * rewrite <- B, <-E. apply rem_loops_preserves_ends.
  Qed.

  Global Instance dec_valid_path r v: dec(valid_path (aut:= aut) r v).
  Proof.
    auto. (*by bounded forall *)
  Qed.

  Lemma bounded_run_is_valid_path r v s s':
    path r v s s' -> path (to_seq(to_bounded ( m:= S (|v|)) r)) v s s'.
  Proof.
    intros V. split.
    - intros k L. cbn in *. rewrite !bounded_unchanged by omega. now apply V.
    - rewrite !bounded_unchanged by comega. apply V.
  Qed.

  Lemma bounded_path_iff_path s s':
    (exists n, n <= Cardinality (state aut) /\ exists (b_r : ConstLengthString (S n)) (b_w: ConstLengthString n), path (to_seq b_r) (mkstring (to_seq b_w) n) s s')
    <-> (exists r v, |v| <= (Cardinality (state aut)) /\ path r v s s').
  Proof.
    split.
    - intros [n [L [R [W P]]]]. exists (to_seq R), (mkstring (to_seq W) n). split; cbn; oauto.
    - intros [r [[w n] [L P]]]. exists n. split; cbn in L.
      + omega.
      + exists (to_bounded r), (to_bounded w). apply bounded_run_is_valid_path; oauto.
        now rewrite <-bounded_unchanged_string.
  Qed.

  Global Instance dec_connected s s' : dec (reachable s s').
  Proof.
    apply (dec_prop_iff (path_length_bounded s s')).
    apply (dec_prop_iff (bounded_path_iff_path s s')).
    auto.
  Qed.

End DecConnected.