From VLSM.Lib Require Import Itauto.
From stdpp Require Import prelude.
From VLSM.Lib Require Import ListExtras.
From VLSM.Core Require Import VLSM PreloadedVLSM VLSMProjections Validator Composition.
From VLSM.Core Require Import ConstrainedVLSM.

Core: State-Annotated VLSMs

This module describes the VLSM obtained by augmenting the states of an existing VLSM with annotations and providing additional validity constraints taking into account the annotations, and a function for updating the annotations following a transition.

Section sec_annotated_vlsm.

Context
  {message : Type}
  (X : VLSM message)
  (annotation : Type)
  .

Record annotated_state : Type := mk_annotated_state
{
  original_state : state X;
  state_annotation : annotation;
}.

Definition annotated_type : VLSMType message :=
  {| label := label X;
     state := annotated_state
  |}.

Context
  (initial_annotation_prop : annotation -> Prop)
  `{Inhabited (sig initial_annotation_prop )}
  .

Definition annotated_initial_state_prop (sa : annotated_state) : Prop :=
  initial_state_prop X (original_state sa) /\ initial_annotation_prop (state_annotation sa).

#[export] Program Instance annotated_initial_state_prop_inhabited :
  Inhabited {sa : annotated_state | annotated_initial_state_prop sa } :=
    populate (exist _ {| original_state := `(vs0 X); state_annotation := `inhabitant; |} _).
Next Obligation.
Proof.
  split; cbn.
  - by destruct (vs0 X).
  - by destruct inhabitant.
Qed.

Context
  (annotated_transition_state :
    label annotated_type -> annotated_state * option message -> annotation)
  .

Definition annotated_valid
  (l : label annotated_type)
  (som : annotated_state * option message)
  : Prop :=
  valid X l (original_state som.1, som.2).

Definition annotated_transition
  (l : label annotated_type)
  (som : annotated_state * option message)
  : annotated_state * option message :=
  let (s', om') := transition X l (original_state som.1, som.2) in
  ({| original_state := s'; state_annotation := annotated_transition_state l som |}, om').

Definition annotated_vlsm_machine : VLSMMachine annotated_type :=
  {| initial_state_prop := fun s : state annotated_type => annotated_initial_state_prop s
  ; initial_message_prop := fun m : message => initial_message_prop X m
  ; valid := annotated_valid
  ; transition := annotated_transition
  |}.

Definition annotated_vlsm : VLSM message := mk_vlsm annotated_vlsm_machine.

Definition annotate_trace_item
  (item : transition_item X)
  (k : annotated_state -> list (transition_item annotated_type))
  (sa : annotated_state)
  : list (transition_item annotated_type) :=
    let sa' :=
      {|
        original_state := destination item;
        state_annotation := annotated_transition_state (l item) (sa , input item );
      |}
    in
      Build_transition_item annotated_type (l item) (input item) sa' (output item) :: k sa'.

Lemma annotate_trace_item_project
  (item : transition_item X)
  (k : annotated_state -> list (transition_item annotated_type))
  (sa : annotated_state)
  : pre_VLSM_embedding_finite_trace_project
      annotated_type X id original_state
      (annotate_trace_item item k sa)
      = item ::
        pre_VLSM_embedding_finite_trace_project
            annotated_type X id original_state
            (k {| original_state := destination item;
                  state_annotation := annotated_transition_state (l item) (sa , input item ) |}).
Proof.
  by destruct item.
Qed.

Definition annotate_trace_from (sa : state annotated_type) (tr : list (transition_item X))
  : list (transition_item annotated_type) :=
  fold_right annotate_trace_item (fun sa => []) tr sa.

Lemma annotate_trace_from_unroll sa item tr
  : annotate_trace_from sa (item :: tr) =
    let sa' :=
      {|
        original_state := destination item;
        state_annotation := annotated_transition_state (l item) (sa , input item )
      |}
    in
      Build_transition_item annotated_type (l item) (input item) sa' (output item)
        ::
      annotate_trace_from sa' tr.
Proof.
  by destruct item.
Qed.

Lemma annotate_trace_from_app sa tr1 tr2
  : annotate_trace_from sa (tr1 ++ tr2) =
    annotate_trace_from sa tr1 ++
      annotate_trace_from (finite_trace_last sa (annotate_trace_from sa tr1)) tr2.
Proof.
  revert sa.
  induction tr1 as [| item tr1]; [done |].
  by intro sa; rewrite <- app_comm_cons, !annotate_trace_from_unroll
  ; simpl; rewrite IHtr1, finite_trace_last_cons.
Qed.

Lemma annotate_trace_from_last_original_state sa sa' tr
  : original_state (finite_trace_last sa (annotate_trace_from sa' tr)) =
    finite_trace_last (original_state sa) tr.
Proof.
  destruct_list_last tr tr' item Heqtr; subst; [done |].
  rewrite annotate_trace_from_app.
  cbn; unfold annotate_trace_item.
  by rewrite! finite_trace_last_is_last.
Qed.

Definition annotate_trace (s : state X) (tr : list (transition_item X))
  : list (transition_item annotated_type) :=
  annotate_trace_from {| original_state := s; state_annotation := ` inhabitant |} tr.

Lemma annotate_trace_last_original_state s s' tr
  : original_state (finite_trace_last s (annotate_trace s' tr)) =
    finite_trace_last (original_state s) tr.
Proof. by apply annotate_trace_from_last_original_state. Qed.

Lemma annotate_trace_project is tr
  : pre_VLSM_embedding_finite_trace_project
      annotated_type X id original_state
      (annotate_trace is tr)
      = tr.
Proof.
  unfold annotate_trace
  ; remember {| original_state := is |} as sa; clear Heqsa; revert sa.
  induction tr as [| item]; [done | intro sa].
  setoid_rewrite annotate_trace_item_project; f_equal.
  by apply IHtr.
Qed.

End sec_annotated_vlsm.

Arguments original_state {_ _ _} _ : assert.
Arguments state_annotation {_ _ _} _ : assert.

Section sec_annotated_vlsm_projections.

Context
  {message : Type}
  (X : VLSM message)
  {annotation : Type}
  (initial_annotation_prop : annotation -> Prop)
  `{Inhabited (sig initial_annotation_prop )}
  (annotated_constraint :
    label (annotated_type X annotation) ->
      annotated_state X annotation * option message -> Prop)
  (annotated_transition_state :
    label (annotated_type X annotation) ->
      annotated_state X annotation * option message -> annotation)
  (AnnotatedX : VLSM message :=
    constrained_vlsm (annotated_vlsm X annotation initial_annotation_prop annotated_transition_state)
      annotated_constraint)
  .

Definition forget_annotations_projection
  : VLSM_embedding AnnotatedX X id original_state.
Proof.
  apply basic_VLSM_strong_embedding; cycle 1; [| by cbv; itauto..].
  intros l [s a] om [s' a'] om'.
  cbn; unfold annotated_transition; cbn
  ; destruct (transition _ _ _) as (_s', _om').
  by inversion 1.
Qed.

End sec_annotated_vlsm_projections.

Section sec_composite_annotated_vlsm_projections.

Specializing projection validator properties to annotated compositions