From Coq Require Import
     Morphisms
.

From ExtLib Require Import
     Structures.Monad.

From ITree Require Import
     Basics.Tacs
     Axioms
     ITree
     Eq
     Basics
.

From Paco Require Import paco.

Import Monads.
Import MonadNotation.
Local Open Scope monad_scope.

Section RuttF.

  Context {E1 E2 : Type → Type}.
  Context {R1 R2 : Type}.
  (* From the point of view of relational parametricity, it would be more fitting
  to replace (REv, RAns) with one REv : ∀ A1 A2, (A1 → A2 → Prop) → (E1 A1 → E2 A2 → Prop).
  Contributions to that effect are welcome. *)
  Context (REv : ∀ (A B : Type), E1 A → E2 B → Prop ).
  Context (RAns : ∀ (A B : Type), E1 A → A → E2 B → B → Prop ).
  Context (RR : R1 → R2 → Prop).
  Arguments REv {A} {B}.
  Arguments RAns {A} {B}.

  Inductive ruttF (sim : itree E1 R1 → itree E2 R2 → Prop) : itree' E1 R1 → itree' E2 R2 → Prop :=
  | EqRet : ∀ (r1 : R1) (r2 : R2),
      RR r1 r2 →
      ruttF sim (RetF r1) (RetF r2)
  | EqTau : ∀ (m1 : itree E1 R1) (m2 : itree E2 R2),
      sim m1 m2 →
      ruttF sim (TauF m1) (TauF m2)
  | EqVis : ∀ (A B : Type) (e1 : E1 A) (e2 : E2 B ) (k1 : A → itree E1 R1) (k2 : B → itree E2 R2),
      REv e1 e2 →
      (∀ (a : A) (b : B), RAns e1 a e2 b → sim (k1 a) (k2 b)) →
      ruttF sim (VisF e1 k1) (VisF e2 k2)
  | EqTauL : ∀ (t1 : itree E1 R1) (ot2 : itree' E2 R2),
      ruttF sim (observe t1) ot2 →
      ruttF sim (TauF t1) ot2
  | EqTauR : ∀ (ot1 : itree' E1 R1) (t2 : itree E2 R2),
      ruttF sim ot1 (observe t2) →
      ruttF sim ot1 (TauF t2).
  Hint Constructors ruttF : itree.

  Definition rutt_ (sim : itree E1 R1 → itree E2 R2 → Prop)
             (t1 : itree E1 R1) (t2 : itree E2 R2) :=
    ruttF sim (observe t1) (observe t2).
  Hint Unfold rutt_ : itree.

  Lemma rutt_monot : monotone2 rutt_.
  Proof.
    red. intros. red; induction IN; eauto with itree.
  Qed.

  Definition rutt : itree E1 R1 → itree E2 R2 → Prop := paco2 rutt_ bot2.

End RuttF.

#[global] Hint Resolve rutt_monot : paco.

Lemma rutt_inv_Tau_l {E1 E2 R1 R2 REv RAns RR} t1 t2 :
  @rutt E1 E2 R1 R2 REv RAns RR (Tau t1) t2 → rutt REv RAns RR t1 t2.
Proof.
  intros. punfold H. red in H. simpl in ×.
  remember (TauF t1) as tt1. genobs t2 ot2.
  hinduction H before t1; intros; try discriminate.
  - inv Heqtt1. pclearbot. pstep. red. simpobs. econstructor; eauto. pstep_reverse.
  - inv Heqtt1. punfold_reverse H.
  - red in IHruttF. pstep. red; simpobs. econstructor; eauto. pstep_reverse.
Qed.

Lemma rutt_add_Tau_l {E1 E2 R1 R2 REv RAns RR} t1 t2 :
  @rutt E1 E2 R1 R2 REv RAns RR t1 t2 → rutt REv RAns RR (Tau t1) t2.
Proof.
  intros. pfold. red. cbn. constructor. pstep_reverse.
Qed.

Lemma rutt_inv_Tau_r {E1 E2 R1 R2 REv RAns RR} t1 t2 :
  @rutt E1 E2 R1 R2 REv RAns RR t1 (Tau t2) → rutt REv RAns RR t1 t2.
Proof.
  intros. punfold H. red in H. simpl in ×.
  pstep. red. remember (TauF t2) as tt2 eqn:Ett2 in H.
  revert t2 Ett2; induction H; try discriminate; intros; inversion Ett2; subst; auto.
  - pclearbot. constructor. pstep_reverse.
  - constructor. eapply IHruttF; eauto.
Qed.

Lemma rutt_add_Tau_r {E1 E2 R1 R2 REv RAns RR} t1 t2 :
  @rutt E1 E2 R1 R2 REv RAns RR t1 t2 → rutt REv RAns RR t1 (Tau t2).
Proof.
  intros. pfold. red. cbn. constructor. pstep_reverse.
Qed.

Lemma rutt_inv_Tau {E1 E2 R1 R2 REv RAns RR} t1 t2 :
  @rutt E1 E2 R1 R2 REv RAns RR (Tau t1) (Tau t2) → rutt REv RAns RR t1 t2.
Proof.
  intros; apply rutt_inv_Tau_r, rutt_inv_Tau_l; assumption.
Qed.

Section euttge_trans_clo.

  Context {E1 E2 : Type → Type} {R1 R2 : Type} (RR : R1 → R2 → Prop).

  (* Closing a relation over itrees under euttge.
     Essentially the same closure as eqit_trans_clo, but heterogeneous
     in the interface argument E.
     We only define the closure under euttge as opposed to eqit_trans_clo
     capturing closure under eq_itree and eutt at the same time, since it's
     the only one we need.
   *)

  (* A transitivity functor *)
  Variant euttge_trans_clo (r : itree E1 R1 → itree E2 R2 → Prop) :
    itree E1 R1 → itree E2 R2 → Prop :=
    eqit_trans_clo_intro t1 t2 t1' t2' RR1 RR2
                         (EQVl: euttge RR1 t1 t1')
                         (EQVr: euttge RR2 t2 t2')
                         (REL: r t1' t2')
                         (LERR1: ∀ x x' y, RR1 x x' → RR x' y → RR x y)
                         (LERR2: ∀ x y y', RR2 y y' → RR x y' → RR x y) :
      euttge_trans_clo r t1 t2.
  Hint Constructors euttge_trans_clo : itree.

  Lemma euttge_trans_clo_mon r1 r2 t1 t2
        (IN : euttge_trans_clo r1 t1 t2)
        (LE : r1 <2= r2) :
    euttge_trans_clo r2 t1 t2.
  Proof.
    destruct IN; econstructor; eauto.
  Qed.

  Hint Resolve euttge_trans_clo_mon : paco.

End euttge_trans_clo.

(*replicate this proof for the models functor*)
(* Validity of the up-to euttge principle *)
Lemma euttge_trans_clo_wcompat E1 E2 R1 R2 (REv : ∀ A B, E1 A → E2 B → Prop)
      (RAns : ∀ A B, E1 A → A → E2 B → B → Prop ) (RR : R1 → R2 → Prop) :
  wcompatible2 (rutt_ REv RAns RR) (euttge_trans_clo RR).
Proof.
  constructor; eauto with paco.
  { red. intros. eapply euttge_trans_clo_mon; eauto. }
  intros.
  destruct PR. punfold EQVl. punfold EQVr. unfold_eqit.
  hinduction REL before r; intros; clear t1' t2'.
  - remember (RetF r1) as x. red.
    hinduction EQVl before r; intros; subst; try inv Heqx; eauto; (try constructor; eauto).
    remember (RetF r3) as x. hinduction EQVr before r; intros; subst; try inv Heqx; (try constructor; eauto).
  - red. remember (TauF m1) as x.
    hinduction EQVl before r; intros; subst; try inv Heqx; try inv CHECK; ( try (constructor; eauto; fail )).
    remember (TauF m3) as y.
    hinduction EQVr before r; intros; subst; try inv Heqy; try inv CHECK; (try (constructor; eauto; fail)).
    pclearbot. constructor. gclo. econstructor; eauto with paco.
  - remember (VisF e1 k1) as x. red.
    hinduction EQVl before r; intros; subst; try discriminate; try (constructor; eauto; fail).
    remember (VisF e2 k3) as y.
    hinduction EQVr before r; intros; subst; try discriminate; try (constructor; eauto; fail).
    dependent destruction Heqx.
    dependent destruction Heqy.
    constructor; auto. intros. apply H0 in H1. pclearbot.
    apply gpaco2_clo.
    econstructor; eauto with itree.
  - remember (TauF t1) as x. red.
    hinduction EQVl before r; intros; subst; try inv Heqx; try inv CHECK; (try (constructor; eauto; fail)).
    pclearbot. punfold REL. constructor. eapply IHREL; eauto.
  - remember (TauF t2) as y. red.
    hinduction EQVr before r; intros; subst; try inv Heqy; try inv CHECK; (try (constructor; eauto; fail)).
    pclearbot. punfold REL. constructor. eapply IHREL; eauto.
Qed.

#[global] Hint Resolve euttge_trans_clo_wcompat : paco.

(* The validity of the up-to euttge entails we can rewrite under euttge
   and hence also eq_itree during coinductive proofs of rutt
*)
#[global] Instance grutt_cong_eqit {R1 R2 : Type} {E1 E2 : Type → Type} {REv : ∀ A B, E1 A → E2 B → Prop}
       {RAns : ∀ A B, E1 A → A → E2 B → B → Prop} {RR1 RR2} {RS : R1 → R2 → Prop} r rg
       (LERR1: ∀ x x' y, (RR1 x x': Prop) → (RS x' y: Prop) → RS x y)
       (LERR2: ∀ x y y', (RR2 y y': Prop) → RS x y' → RS x y) :
  Proper (eq_itree RR1 ==> eq_itree RR2 ==> flip impl)
         (gpaco2 (rutt_ REv RAns RS) (euttge_trans_clo RS) r rg).
Proof.
  repeat intro. gclo. econstructor; eauto;
    try eapply eqit_mon; try apply H; try apply H0; auto.
Qed.

Global Instance grutt_cong_euttge {R1 R2 : Type} {E1 E2 : Type → Type} {REv : ∀ A B, E1 A → E2 B → Prop}
       {RAns : ∀ A B, E1 A → A → E2 B → B → Prop} {RR1 RR2} {RS : R1 → R2 → Prop} r rg
       (LERR1: ∀ x x' y, (RR1 x x': Prop) → (RS x' y: Prop) → RS x y)
       (LERR2: ∀ x y y', (RR2 y y': Prop) → RS x y' → RS x y) :
  Proper (euttge RR1 ==> euttge RR2 ==> flip impl)
         (gpaco2 (rutt_ REv RAns RS) (euttge_trans_clo RS) r rg).
Proof.
  repeat intro. gclo. econstructor; eauto.
Qed.