Inductive trace {E : Type → Type} {R : Type} : Type :=
| TEnd : trace
| TRet : R → trace
| TEventEnd : ∀ {X}, E X → trace
| TEventResponse : ∀ {X}, E X → X → trace → trace
.

(* Append two traces, ignoring the end of the first trace. *)
Fixpoint app_trace {E R S} (tr1 : @trace E R) (tr2 : @trace E S) : @trace E S :=
  match tr1 with
  | TEventResponse e x tr ⇒ TEventResponse e x (app_trace tr tr2)
  | _ ⇒ tr2
  end.

(* Get the value in the TRet at the end of the trace, if it exists. *)
Fixpoint trace_ret {E R} (tr : @trace E R) : option R :=
  match tr with
  | TRet r ⇒ (Some r)
  | TEventResponse _ _ tr ⇒ trace_ret tr
  | _ ⇒ None
  end.

Inductive is_traceF {E : Type → Type} {R : Type} :
  itreeF E R (itree E R) → @trace E R → Prop :=
| TraceEmpty : ∀ t, is_traceF t TEnd
| TraceRet : ∀ r, is_traceF (RetF r) (TRet r)
| TraceTau : ∀ t tr,
    is_traceF (observe t) tr →
    is_traceF (TauF t) tr
| TraceVisEnd : ∀ X (e : E X) k,
    is_traceF (VisF e k) (TEventEnd e)
| TraceVisContinue : ∀ X (e : E X) (x : X) k tr,
    is_traceF (observe (k x)) tr →
    is_traceF (VisF e k) (TEventResponse e x tr)
.

Definition is_trace {E : Type → Type} {R : Type} (t : itree E R) :=
  is_traceF (observe t).

(* t1 ⊑ t2 *)
Definition trace_incl {E : Type → Type} {R : Type} :
  itree E R → itree E R → Prop :=
  fun t1 t2 ⇒
    ∀ tr, is_trace t1 tr → is_trace t2 tr.

(* t1 ≡ t2 *)
Definition trace_eq {E : Type → Type} {R : Type} :
  itree E R → itree E R → Prop :=
  fun t1 t2 ⇒
    trace_incl t1 t2 ∧ trace_incl t2 t1.

Lemma is_traceF_tau : ∀ {E R} (t : itree E R) tr,
    is_traceF (observe t) tr ↔
    is_traceF (TauF t) tr.
Proof.
  intros. split; intros.
  - constructor. remember (observe t).
    generalize dependent t.
    induction H; intros; subst; constructor; eapply IHis_traceF; auto.
  - inversion H; subst; try constructor; auto.
Qed.

Lemma sutt_trace_incl : ∀ {E R} (t1 t2 : itree E R),
    sutt eq t1 t2 → trace_incl t1 t2.
Proof.
  red. intros. red in H0. remember (observe t1).
  generalize dependent t1. generalize dependent t2.
  induction H0; intros; try solve [constructor].
  - punfold H. rewrite <- Heqi in H.
    remember (RetF _). remember (observe t2).
    generalize dependent t2.
    induction H; intros; try inv Heqi0; red; rewrite <- Heqi1; constructor.
    eapply IHsuttF; eauto.
  - apply IHis_traceF with (t1:=t); auto.
    apply sutt_inv_tau_left. red. red in H. rewrite <- Heqi in H. auto.
  - punfold H. rewrite <- Heqi in H.
    remember (VisF _ _). remember (observe t2).
    generalize dependent t2.
    induction H; intros; try discriminate.
    + inv_Vis. subst. red. rewrite <- Heqi1. constructor.
    + red. rewrite <- Heqi1. constructor. eapply IHsuttF; eauto.
  - punfold H. rewrite <- Heqi in H.
    remember (VisF _ _). remember (observe t2).
    generalize dependent t2.
    induction H; intros; try discriminate.
    + inv_Vis. pclearbot. subst. red. rewrite <- Heqi1. constructor.
      eapply IHis_traceF; auto with itree.
    + red. rewrite <- Heqi1. constructor. apply IHsuttF; auto.
Qed.

Lemma eutt_trace_eq : ∀ {E R} (t1 t2 : itree E R),
    t1 ≈ t2 → trace_eq t1 t2.
Proof.
  split.
  - apply sutt_trace_incl; apply eutt_sutt; auto.
  - symmetry in H. apply sutt_trace_incl; apply eutt_sutt; auto.
Qed.

Lemma eq_trace_inv {E R} (t1 t2 : @trace E R) (H : t1 = t2)
  : match t1, t2 with
    | TEnd, TEnd ⇒ True
    | TRet r1, TRet r2 ⇒ r1 = r2
    | TEventEnd e1, TEventEnd e2 ⇒ ∃ p, eqeq E p e1 e2
    | TEventResponse e1 x1 t1, TEventResponse e2 x2 t2 ⇒
      ∃ p, eqeq E p e1 e2 ∧ eqeq (fun X ⇒ X) p x1 x2 ∧ t1 = t2
    | _, _ ⇒ False
    end.
Proof.
  destruct H, t1; auto.
  - ∃ eq_refl; cbn; auto.
  - ∃ eq_refl; cbn; auto.
Qed.

Lemma trace_incl_sutt : ∀ {E R} (t1 t2 : itree E R),
    trace_incl t1 t2 → sutt eq t1 t2.
Proof.
  intros E R. pcofix CIH. pstep. intros t1 t2 Hincl.
  unfold trace_incl in ×. unfold is_trace in ×.
  destruct (observe t1).
  - assert (H : is_traceF (RetF r0 : itreeF E R (itree E R)) (TRet r0)) by constructor.
    apply Hincl in H. clear Hincl. destruct (observe t2); inv H.
    + constructor. auto.
    + constructor.
      remember (observe t). remember (TRet _).
      generalize dependent t.
      induction H1; intros; try inv Heqt0; auto with itree.
      constructor. eapply IHis_traceF; eauto.
  - constructor. right. apply CIH. intros. apply Hincl. constructor. auto.
  - assert (H: is_traceF (VisF e k) (TEventEnd e)) by constructor.
    apply Hincl in H. destruct (observe t2); inv H.
    + constructor.
      assert (∀ tr, is_traceF (VisF e k) tr → is_traceF (observe t) tr).
      {
        intros. rewrite is_traceF_tau. apply Hincl; auto.
      }
      clear Hincl. rename H into Hincl.
      remember (observe t). remember (TEventEnd _).
      generalize dependent t.
      induction H1; intros; try discriminate.
      × constructor. eapply IHis_traceF; eauto.
        intros. rewrite is_traceF_tau. apply Hincl; auto.
      × apply eq_trace_inv in Heqt0; destruct Heqt0 as [<- <-].
        subst. constructor. intro. right. apply CIH. intros.
        assert (is_traceF (VisF e k) (TEventResponse e x tr)) by (constructor; auto).
        apply Hincl in H1. inv H1. ddestruction. auto.
    + ddestruction. constructor. intro. right. apply CIH. intros.
      assert (is_traceF (VisF e0 k) (TEventResponse e0 x tr)) by (constructor; auto).
      apply Hincl in H0. inv H0. ddestruction; auto.
Qed.

Theorem trace_incl_iff_sutt : ∀ {E R} (t1 t2 : itree E R),
    sutt eq t1 t2 ↔ trace_incl t1 t2.
Proof.
  split.
  - apply sutt_trace_incl.
  - apply trace_incl_sutt.
Qed.

Lemma trace_eq_eutt : ∀ {E R} (t1 t2 : itree E R),
    trace_eq t1 t2 → t1 ≈ t2.
Proof.
  intros E R t1 t2 [? ?]. apply sutt_eutt.
  - apply trace_incl_sutt; auto.
  - apply trace_incl_sutt in H0. clear H.
    generalize dependent t1. generalize dependent t2. pcofix CIH; pstep; intros.
    punfold H0. induction H0; constructor; try red; pclearbot; eauto with paco itree.
    right. rewrite itree_eta'. eauto with paco itree.
Qed.

Theorem trace_eq_iff_eutt : ∀ {E R} (t1 t2 : itree E R),
    t1 ≈ t2 ↔ trace_eq t1 t2.
Proof.
  split.
  - apply eutt_trace_eq.
  - apply trace_eq_eutt.
Qed.

Inductive event (E : Type → Type) : Type :=
| Event : ∀ {X}, E X → X → event E
.

(* step_ ev t' t if t steps to t' (read right to left!)
   with visible event ev. *)
Inductive step_ {E : Type → Type} {R : Type}
          (ev : event E) (t' : itree E R) :
  itree E R → Prop :=
| StepTau : ∀ t, step_ ev t' t → step_ ev t' (Tau t)
| StepVis : ∀ X (e : E X) (x : X) k,
    ev = Event _ e x →
    t' = k x →
    step_ ev t' (Vis e k)
.

Definition step {E : Type → Type} {R : Type}
           (ev : event E) (t t' : itree E R) : Prop :=
  step_ ev t' t.

CoInductive simulates {E : Type → Type} {R : Type} :
  itree E R → itree E R → Prop :=
| SimStep : ∀ t1 t2,
    (∀ ev t1',
        step ev t1 t1' →
        ∃ t2', step ev t2 t2' ∧ simulates t1' t2') →
    simulates t1 t2.

Theorem simulates_trace_incl {E : Type → Type} {R : Type} :
  ∀ t1 t2 : itree E R,
    simulates t1 t2 → trace_incl t1 t2.
Proof.
Abort.

(* Set-of-traces monad *)
Definition traces (E : Type → Type) (R : Type) : Type :=
  @trace E R → Prop.

Definition bind_traces {E : Type → Type} {R S : Type}
           (ts : traces E R) (kts : R → traces E S) : traces E S :=
  fun tr ⇒
    (tr = TEnd ∧ ts TEnd) ∨
    (∃ X (e : E X), tr = TEventEnd e ∧ ts (TEventEnd e)) ∨
    (∃ r tr1 tr2,
        tr = app_trace tr1 tr2 ∧
        trace_ret tr1 = Some r ∧
        ts tr1 ∧
        kts r tr).

Definition ret_traces {E : Type → Type} {R : Type} :
  R → traces E R :=
  fun r tr ⇒
    tr = TEnd ∨ tr = TRet r.