From Stdlib Require Import List Program Lia ConstructiveEpsilon Arith Classes.Morphisms.
From mathcomp Require Import zify.
From Equations Require Import Equations.
Require Import extra_principles Util continuity_zoo_Prop partial embed_nat Common Modulus Neighborhood.
Import EmbedNatNotations PartialTactics ListNotations.

Set Bullet Behavior "Strict Subproofs".

Section PartialZoo.

Context `{Partiality : partiality}.

Definition ext_tree {Q A R : Type} := list A -> part ( @result Q R ).

Fixpoint eval_ext_tree_aux [Q A R : Type] (tau : ext_tree (R := R)) (f : Q -> part A)
  (n : nat) (l : list A) : part result
:=
  match n with
  | S k => bind (tau l) (fun r => match r with
    | ask q => bind (f q) (fun a => eval_ext_tree_aux tau f k (a :: l))
    | output r => ret (output r)
    end)
  | 0 => tau l
  end.

Add Parametric Morphism Q A R : ( @eval_ext_tree_aux Q A R )
  with signature pointwise_relation _ ( @equiv _ result )
    ==> pointwise_relation _ ( @equiv _ A )
    ==> eq
    ==> eq
    ==> @equiv _ result
  as eval_ext_tree_aux_morph.
Proof.
  intros tau1 tau2 Htau f1 f2 Hf n.
  induction n; intros l.
  - psimpl.
  - cbn.
    rewrite Htau.
    apply bind_morph; try reflexivity.
    intros r.
    destruct r; psimpl.
    rewrite Hf.
    setoid_rewrite IHn.
    reflexivity.
Qed.

Definition eval_ext_tree [Q A R : Type] (tau : ext_tree (R := R)) (f : Q -> part A) (n : nat) :
  part result
:=
  eval_ext_tree_aux tau f n nil.

Definition tree_fun_for [Q A R : Type] (tau : @ext_tree Q A R) (F : (Q -> part A) -> part R) :=
  forall (f : Q -> part A) (r : R),
    F f =! r <-> exists (n : nat), eval_ext_tree tau f n =! output r.

Definition well_founded_ext_tree [Q A R : Type] (tau : @ext_tree Q A R) : Prop :=
  (forall (l l' : list A), ter (tau (l' ++ l)) -> ter (tau l)) /\
  (forall (r : R) (l l' : list A), tau l =! output r -> tau (l' ++ l) =! output r).

Definition well_founded_tree_fun_for [Q A R : Type] (tau : @ext_tree Q A R) (F : (Q -> part A) -> part R) : Prop :=
  well_founded_ext_tree tau /\ tree_fun_for tau F.

Definition well_founded_tree_fun_cont [Q A R : Type] (F : (Q -> part A) -> part R) : Prop :=
  exists (tau : ext_tree), well_founded_tree_fun_for tau F.

Definition strict_ext_tree [Q A R : Type] (tau : @ext_tree Q A R) : Prop :=
  forall (l : list A), ter (tau l) <-> (exists (r : R) (l' : list A),
    tau (l' ++ l) =! output r).

Definition strict_well_founded_tree_fun_for [Q A R : Type] (tau : @ext_tree Q A R) (F : (Q -> part A) -> part R) : Prop :=
  strict_ext_tree tau /\ well_founded_tree_fun_for tau F.

Definition strict_well_founded_tree_fun_cont [Q A R : Type] (F : (Q -> part A) -> part R) : Prop :=
  exists (tau : ext_tree), strict_well_founded_tree_fun_for tau F.

Definition tree_fun_cont [Q A R : Type] (F : (Q -> part A) -> part R) : Prop :=
  exists (tau : ext_tree), tree_fun_for tau F.

Lemma eval_ext_tree_aux_det [Q A R : Type] (tau : ext_tree (R := R)) (f : Q -> part A) :
  forall (n m : nat) (l : list A) (r r' : R),
    eval_ext_tree_aux tau f n l =! output r ->
    eval_ext_tree_aux tau f m l =! output r' ->
    r = r'.
Proof.
  enough (forall (n m : nat) (l : list A) (r r' : R),
    n <= m ->
    eval_ext_tree_aux tau f n l =! output r ->
    eval_ext_tree_aux tau f m l =! output r' ->
    r = r'). {
    intros n m l r r' Hr Hr'.
    destruct (Nat.le_ge_cases n m) as [ H' | H' ];
    [ | symmetry ]; eapply H; eassumption.
  }
  intros n m l r r' H.
  clear H.
  revert m l.
  induction n as [| n IHn]; intros m; destruct m as [| m];
  intros l Hr Hr'; simpl in Hr, Hr'.
  - assert (E : output r = output r') by (eapply hasvalue_det; eassumption).
    inversion E; auto.
  - psimpl.
    assert (x = output r) by (eapply hasvalue_det; eassumption).
    subst.
    eapply ret_hasvalue_inv in H0.
    inversion H0; auto.
  - psimpl.
    assert (output r' = x) by (eapply hasvalue_det; eassumption).
    subst.
    eapply ret_hasvalue_inv in H0.
    inversion H0; auto.
  - psimpl.
    assert (x0 = x) by (eapply hasvalue_det; eassumption).
    subst.
    destruct x.
    + psimpl.
      assert (x0 = x) by (eapply hasvalue_det; eassumption).
      subst.
      eapply IHn; eassumption.
    + assert (E : output r = output r') by (eapply hasvalue_det; eassumption).
      inversion E; auto.
Qed.

Lemma eval_ext_tree_aux_succ_ter [Q A R : Type] (tau : ext_tree (R := R)) (f : Q -> part A) :
  forall n l, ter (eval_ext_tree_aux tau f (S n) l) ->
  ter (eval_ext_tree_aux tau f n l).
Proof.
  intros n.
  induction n; cbn; intros l Hf;
  destruct Hf as (o & Ho); psimpl.
  - eexists; eauto.
  - destruct x; try now (eexists; psimpl; eauto).
    psimpl.
    assert (Hn : ter (eval_ext_tree_aux tau f (S n) (x :: l)))
    by (cbn; eexists; psimpl).
    destruct (IHn _ Hn) as (o' & Ho').
    exists o'.
    psimpl.
Qed.

Lemma eval_ext_tree_aux_succ_result [Q A R : Type] (tau : ext_tree (R := R)) (f : Q -> part A) :
  forall n l r, eval_ext_tree_aux tau f n l =! output r ->
  eval_ext_tree_aux tau f (S n) l =! output r.
Proof.
  intros n.
  induction n; intros l r Hf; cbn in Hf.
  - psimpl.
  - apply bind_hasvalue in Hf as ([q | r'] & Ho & Ho').
    + apply bind_hasvalue in Ho' as (a & Ha & Ha').
      apply IHn in Ha'.
      cbn in *.
      psimpl.
    + apply ret_hasvalue_inv in Ho'; inversion Ho'; subst.
      psimpl.
Qed.

Fixpoint tree_fun_trace_aux [Q A R : Type] (tau : @ext_tree Q A R)
  (f : Q -> part A) (n : nat) (l : list A) : part (list Q)
:=
  match n with
  | O => ret nil
  | S n => bind (tau l) (fun o => match o with
    | output _ => ret nil
    | ask q => let p := bind (f q)
        (fun a => tree_fun_trace_aux tau f n (a :: l))
      in bind p (fun l => ret (q :: l))
    end)
  end.

Definition tree_fun_trace [Q A R : Type] (tau : @ext_tree Q A R)
  (f : Q -> part A) (n : nat) : part (list Q)
:= tree_fun_trace_aux tau f n nil.

Lemma tree_fun_trace_aux_det [Q A R : Type] :
  forall (tau : @ext_tree Q A R) (f : Q -> part A) (n m : nat)
    (r : R) (l : list Q) (l' : list A),
  eval_ext_tree_aux tau f n l' =! output r ->
  eval_ext_tree_aux tau f m l' =! output r ->
  tree_fun_trace_aux tau f n l' =! l ->
  tree_fun_trace_aux tau f m l' =! l.
Proof.
  intros tau f n.
  induction n as [| n IHn ];
  intros m; destruct m as [| m ];
  intros r l l' Hn Hm Hn'; simpl in *.
  - assumption.
  - psimpl.
    assert (x = output r) by (eapply hasvalue_det; eauto).
    subst.
    psimpl.
  - psimpl.
    assert (x = output r) by (eapply hasvalue_det; eauto).
    assert (x0 = output r) by (eapply hasvalue_det; eauto).
    subst.
    assumption.
  - psimpl.
    assert (x0 = x) by (eapply hasvalue_det; eauto).
    assert (x1 = x) by (eapply hasvalue_det; eauto).
    subst.
    destruct x as [q | ?].
    2: assumption.
    psimpl.
    assert (x0 = x) by (eapply hasvalue_det; eauto).
    assert (x2 = x) by (eapply hasvalue_det; eauto).
    subst.
    eapply IHn; eauto.
Qed.

Lemma tree_fun_trace_ter [Q A R : Type] :
  forall (tau : @ext_tree Q A R) (f : Q -> part A) (n : nat) (l' : list A),
  ter (eval_ext_tree_aux tau f n l') ->
  ter (tree_fun_trace_aux tau f n l').
Proof.
  intros f tau n.
  induction n as [|n IHn];
  intros l' (r & Hr); cbn in Hr; psimpl.
  1: eexists; psimpl.
  destruct x as [q | r'].
  2: eexists; psimpl.
  psimpl.
  eapply bind_ter; try eassumption.
  simpl.
  eapply bind_ter'.
  2: intros; eexists; psimpl.
  eapply bind_ter; try eassumption.
  eapply IHn.
  eexists; eassumption.
Qed.

Lemma tree_fun_trace_val [Q A R : Type] :
  forall (tau : @ext_tree Q A R) f n l', equiv (eval_ext_tree_aux tau f n l')
    (bind (tree_fun_trace_aux tau f n l')
      (fun l => bind (fold_part_list (map f l))
        (fun l => tau (l ++ l')))).
Proof.
  intros tau f n.
  induction n; intros l; cbn.
  - rewrite equiv_bind_ret.
    cbn.
    rewrite equiv_bind_ret.
    reflexivity.
  - rewrite equiv_bind_mult.
    eapply equiv_bind; try reflexivity.
    intros [q | r] Htau _.
    + repeat rewrite equiv_bind_mult.
      setoid_rewrite equiv_bind_ret.
      eapply equiv_bind; try reflexivity.
      intros a Ha _.
      rewrite IHn.
      eapply equiv_bind; try reflexivity.
      intros l' _ _.
      change (map f (q :: l')) with (f q :: map f l').
      rewrite fold_part_list_cons.
      rewrite equiv_bind_mult.
      erewrite (equiv_bind_hasvalue _ _ _ Ha).
      rewrite equiv_bind_mult.
      setoid_rewrite equiv_bind_ret.
      eapply equiv_bind; try reflexivity.
      intros ? _ _.
      rewrite <- app_assoc.
      reflexivity.
    + rewrite equiv_bind_ret.
      cbn.
      rewrite equiv_bind_ret.
      eapply equiv_hasvalue; psimpl.
Qed.

Lemma eval_ext_tree_aux_succ [Q A R : Type] (tau : @ext_tree Q A R) (f : Q -> part A) :
  forall (n : nat) (l : list A),
    equiv (eval_ext_tree_aux tau f (S n) l)
      (bind (eval_ext_tree_aux tau f n l) (fun res => match res with
      | output r => ret (output r)
      | ask q => bind (f q) (fun a => bind (tree_fun_trace_aux tau f n l)
          (fun lq => bind (fold_part_list (map f lq))
            (fun la => tau (a :: la ++ l)))
        )
      end)).
Proof with try reflexivity.
  intros n.
  induction n as [|n IHn]; intros l. {
    cbn.
    apply equiv_bind...
    intros res Hres _.
    destruct res...
    apply equiv_bind...
    intros a Ha _.
    rewrite equiv_bind_ret.
    cbn.
    rewrite equiv_bind_ret...
  }
  cbn in *.
  setoid_rewrite equiv_bind_mult.
  apply equiv_bind...
  intros res Hres _.
  destruct res.
  2: eapply equiv_hasvalue; psimpl.
  setoid_rewrite equiv_bind_mult.
  apply equiv_bind...
  intros a Ha _.
  rewrite IHn.
  apply equiv_bind...
  intros res' Hres' _.
  destruct res'; psimpl.
  apply equiv_bind...
  intros a0 Ha0 _.
  rewrite equiv_bind_hasvalue with (p := tau l); eauto.
  simpl.
  repeat setoid_rewrite equiv_bind_mult.
  rewrite equiv_bind_hasvalue with (p := f q); eauto.
  apply equiv_bind...
  intros lq Hlq _.
  rewrite equiv_bind_ret.
  simpl.
  rewrite fold_part_list_cons.
  repeat setoid_rewrite equiv_bind_mult.
  rewrite equiv_bind_hasvalue with (p := f q); eauto.
  setoid_rewrite equiv_bind_ret.
  apply equiv_bind...
  intros la Hla _.
  rewrite <- app_assoc...
Qed.

Lemma eval_ext_tree_succ [Q A R : Type] (tau : @ext_tree Q A R) (f : Q -> part A) :
  forall (n : nat),
    equiv (eval_ext_tree tau f (S n))
      (bind (eval_ext_tree tau f n) (fun res => match res with
      | output r => ret (output r)
      | ask q => bind (f q) (fun a => bind (tree_fun_trace tau f n)
          (fun lq => bind (fold_part_list (map f lq))
            (fun la => tau (a :: la)))
        )
      end)).
Proof with try reflexivity.
  intros n.
  unfold eval_ext_tree.
  rewrite (eval_ext_tree_aux_succ tau f n nil).
  apply equiv_bind...
  intros res Hres _.
  destruct res...
  setoid_rewrite app_nil_r with (l := _ :: _)...
Qed.

Lemma tree_fun_trace_aux_succ [Q A R : Type] (tau : ext_tree (R := R)) (f : Q -> part A) :
  forall (n : nat) (la : list A),
    equiv (tree_fun_trace_aux tau f (S n) la)
      (bind (tree_fun_trace_aux tau f n la) (fun lq => bind (fold_part_list (map f lq))
        (fun la' => bind (tau (la' ++ la)) (fun res => match res with
        | output _ => ret lq
        | ask q => bind (f q) (fun _ => ret (lq ++ [ q ]))
        end))
      )).
Proof with try reflexivity.
  intros n.
  induction n as [|n IHn]; intros la. {
    cbn.
    rewrite equiv_bind_ret.
    cbn.
    rewrite equiv_bind_ret.
    cbn.
    eapply equiv_bind...
    intros res Hres _.
    destruct res...
    rewrite equiv_bind_mult.
    setoid_rewrite equiv_bind_ret...
  }
  cbn in *.
  rewrite equiv_bind_mult.
  eapply equiv_bind...
  intros res Hres _.
  destruct res.
  2:eapply equiv_hasvalue; psimpl.
  repeat setoid_rewrite equiv_bind_mult.
  eapply equiv_bind...
  intros a Ha _.
  setoid_rewrite equiv_bind_ret.
  specialize (IHn (a :: la)).
  intros lq; split; intros Hl.
  - simpl_assms.
    eassert (bind (tau (a :: la)) _ =! _) as Hn%IHn by (rewrite bind_hasvalue_given; eauto).
    simpl_assms.
    change (map f (q :: ?a0)) with (f q :: map f a0).
    setoid_rewrite fold_part_list_cons.
    psimpl.
    1: rewrite <- app_assoc; eauto.
    destruct x3; psimpl.
  - change (map f (q :: ?x)) with (f q :: map f x) in Hl.
    setoid_rewrite fold_part_list_cons in Hl.
    simpl_assms.
    rewrite <- app_assoc in H1.
    simpl in H1.
    assert (x2 = a) by (eapply hasvalue_det; eauto); subst.
    eassert (Hn : bind _ _ =! match x1 with
    | ask q0 => x ++ [ q0 ]
    | output _ => x
    end) by (apply IHn; destruct x1; psimpl).
    psimpl.
    destruct x0.
    + psimpl.
      destruct x1; psimpl.
      clear H9 H5.
      destruct x; inversion H10; subst; psimpl.
    + destruct x1; psimpl.
Qed.

Lemma tree_fun_trace_succ [Q A R : Type] (tau : ext_tree (R := R)) (f : Q -> part A) :
  forall (n : nat), equiv (tree_fun_trace tau f (S n))
      (bind (tree_fun_trace tau f n) (fun lq => bind (fold_part_list (map f lq))
        (fun la => bind (tau la) (fun res => match res with
        | output _ => ret lq
        | ask q => bind (f q) (fun _ => ret (lq ++ [ q ]))
        end))
      )).
Proof.
  intros n.
  rewrite (tree_fun_trace_aux_succ tau f n nil).
  setoid_rewrite app_nil_r.
  reflexivity.
Qed.

Lemma eval_ext_tree_aux_ask_succ [Q A R : Type] (tau : ext_tree (R := R)) (f : Q -> part A) :
  forall (n : nat) l' q a l,
  tree_fun_trace_aux tau f n l' =! l ->
  eval_ext_tree_aux tau f n l' =! ask q -> f q =! a ->
  exists lf, fold_part_list (map f l) =! lf /\
    equiv (tau (a :: lf ++ l'))
      (eval_ext_tree_aux tau f (S n) l').
Proof.
  intros n.
  induction n; intros l' q a l H1l' H2l' Ha; cbn in H1l', H2l'.
  - apply ret_hasvalue_inv in H1l'.
    subst.
    exists nil; split; psimpl.
    erewrite equiv_bind_hasvalue; eauto.
    cbn.
    erewrite equiv_bind_hasvalue; eauto.
    reflexivity.
  - psimpl.
    assert (x0 = x) by (eapply hasvalue_det; eauto); subst.
    destruct x.
    + rewrite bind_hasvalue in H0, H2.
      destruct H0 as (a' & Ha' & H').
      destruct H2 as (l1 & Hl1 & Hl1').
      apply bind_hasvalue in Hl1 as (a1 & Ha1' & H1').
      assert (a1 = a') by (eapply hasvalue_det; eauto); subst.
      apply ret_hasvalue_inv in Hl1'.
      subst.
      specialize (IHn _ _ _ _ H1' H' Ha).
      destruct IHn as (lf & Hlf & Hlf').
      exists (lf ++ [a']); split.
      * change (map f (q0 :: l1)) with (f q0 :: map f l1).
        rewrite fold_part_list_cons.
        rewrite equiv_bind_hasvalue; eauto.
        rewrite equiv_bind_hasvalue; eauto.
        psimpl.
      * rewrite <- app_assoc.
        simpl.
        rewrite equiv_bind_hasvalue; eauto.
        rewrite equiv_bind_hasvalue; eauto.
    + apply ret_hasvalue_inv in H0.
      inversion H0.
Qed.

Lemma tree_fun_trace_aux_mod [Q A R : Type] :
  forall (tau : @ext_tree Q A R) (f g : Q -> part A) (r : R) (n : nat)
    (l : list Q) (l' : list A),
  eval_ext_tree_aux tau f n l' =! output r ->
  tree_fun_trace_aux tau f n l' =! l ->
  Forall (fun q => equiv (f q) (g q)) l ->
  tree_fun_trace_aux tau g n l' =! l.
Proof.
  intros tau f g r n.
  induction n as [|n IHn]; intros l l' Hf Hf' Hg; simpl in *.
  1: assumption.
  psimpl.
  assert (x0 = x) by (eapply hasvalue_det; eauto).
  subst.
  destruct x.
  2: assumption.
  psimpl.
  - inversion Hg.
    apply H7.
    apply H2.
  - assert (x1 = x) by (eapply hasvalue_det; eauto).
    subst.
    inversion Hg.
    apply IHn; eauto.
Qed.

Lemma tree_fun_trace_mod [Q A R : Type] :
  forall (tau : @ext_tree Q A R) (n : nat) (l : list A),
  modulus (fun f => eval_ext_tree_aux tau f n l)
    (fun f => tree_fun_trace_aux tau f n l).
Proof.
  intros tau n l f Hter.
  apply tree_fun_trace_ter in Hter as Hter'.
  destruct Hter' as (l' & Hl').
  exists l'.
  split; try assumption.
  intros _ g Hg.
  generalize dependent l.
  generalize dependent l'.
  induction n as [|n IHn]; intros l' Hg l Hter Hl'; psimpl.
  simpl in Hl'.
  psimpl.
  eapply equiv_bind'; try eassumption.
  destruct x as [q | r]; try psimpl.
  eapply equiv_bind'.
  - eassumption.
  - inversion Hg. subst. apply H4. assumption.
  - destruct Hter as (x' & Hx').
    simpl in Hx'.
    apply bind_hasvalue in Hx' as (o & Ho & Ho').
    assert (o = ask q) by (eapply hasvalue_det; eassumption).
    subst.
    psimpl.
    eapply IHn with (l' := x) (l := x0 :: l).
    + inversion Hg; assumption.
    + assert (x0 = x1) by (eapply hasvalue_det; eassumption).
      subst.
      eexists; eassumption.
    + assumption.
Qed.

Lemma tree_fun_trace_ter_fold [Q A R : Type] (tau : @ext_tree Q A R) (f : Q -> part A) :
  forall (n : nat) (l' : list A) (l : list Q),
    ter (eval_ext_tree_aux tau f n l') ->
    tree_fun_trace_aux tau f n l' =! l -> ter (fold_part_list (map f l)).
Proof.
  intros n l' l (o & Ho) Hn.
  apply tree_fun_trace_val in Ho.
  rewrite bind_hasvalue_given in Ho; psimpl.
  eexists; psimpl.
Qed.

Definition tree_fun_cont_mod_fn [Q A R : Type]
  `{D : EqDec R} (F : (Q -> part A) -> part R)
  (tau : list A -> part ( @result Q R )) (Htau : tree_fun_for tau F)
  : (Q -> part A) -> part (list Q).
Proof.
  set (P f (H : ter (F f)) q :=
    let (n, m) := unembed q in
    seval (eval_ext_tree tau f n) m = Some (output (eval H))
  ).
  assert (DPr : forall f Hter k, {P f Hter k} + {~P f Hter k}). {
    intros f Hter k.
    unfold P.
    destruct (unembed k) as (n, m) eqn:E.
    destruct (seval (eval_ext_tree tau f n) m).
    2: auto.
    destruct r.
    1: auto.
    destruct (eq_dec r (eval Hter)).
    - subst; auto.
    - right.
      intros H.
      inversion H.
      auto.
  }
  assert (EPr : forall f Hter, exists k, P f Hter k). {
    intros f (r & Hf).
    pose proof (Hf' := Hf).
    apply Htau in Hf' as (n & Hn).
    apply seval_hasvalue in Hn as (m & Hm).
    exists (embed (n, m)).
    unfold P.
    rewrite embedP.
    eenough (r = eval _) as <- by apply Hm.
    eapply hasvalue_det; eauto.
    apply eval_hasvalue.
  }
  unshelve eset (N := _ : forall (f : Q -> part A) (Hter : ter (F f)),
    { k : nat | P f Hter k }). {
    intros f Hter.
    apply constructive_indefinite_ground_description_nat;
    auto.
  }
  set (M f := bind (proof_ter (F f)) (fun Hter =>
    let (n, _) := unembed (proj1_sig (N f Hter)) in
    tree_fun_trace tau f n
  )).
  exact M.
Defined.

Lemma tree_fun_cont_modulus [Q A R : Type] `{D : EqDec R}
  (F : (Q -> part A) -> part R)
  (tau : list A -> part ( @result Q R ))
  (Htau : tree_fun_for tau F)
  : modulus F (tree_fun_cont_mod_fn F tau Htau).
Proof.
  intros f Hf.
  pose proof (Hf' := Hf).
  destruct Hf' as (r & Hr).
  unfold tree_fun_cont_mod_fn.
  assert (Hf' : ter (proof_ter (F f))) by (apply proof_ter_spec; assumption).
  destruct Hf' as (Hter & Hter').
  assert (forall {A : Type} (P Q : A -> Prop),
    (forall x, P x -> Q x) -> (exists x, P x) -> exists x, Q x) by firstorder.
  eapply H with (P := fun l =>
    (let (n, _) := unembed (` _) in
    tree_fun_trace tau f n) =! l /\ modulus_at F f l
  ). {
    intros l (Hl & ?).
    intros .
    split.
    2: auto.
    eapply bind_hasvalue.
    exists Hter.
    split; auto.
    eassumption.
  }
  eset (c := constructive_indefinite_ground_description_nat _ _ _).
  set (P := (fun! m, n0 =>
    seval (A:=result) (eval_ext_tree tau f m) n0 =
    Some (output (eval (x:=F f) Hter)))).
  fold P in c.
  destruct c as (k & Hk).
  destruct (unembed k) as (m & n) eqn:E.
  simpl.
  unfold P in Hk.
  rewrite E in Hk |- *.
  assert (Hf' : hasvalue (eval_ext_tree tau f m) (output (eval Hter)))
  by (rewrite seval_hasvalue; eauto).
  clear E k Hk n Hter' P Hf H.
  destruct (tree_fun_trace_ter tau f m nil) as (l & Hl).
  1: eexists; eassumption.
  exists l.
  split; try assumption.
  intros _ g Hg.
  assert (Hfg : equiv (eval_ext_tree tau f m) (eval_ext_tree tau g m)). {
    destruct (tree_fun_trace_mod tau m nil f (ltac:(eexists; eauto)))
    as (l' & Hl'' & Hl').
    assert (l = l') by (eapply hasvalue_det; eassumption); subst.
    clear Hl''.
    specialize (Hl' (ltac:(eexists; eauto)) g Hg).
    eapply Hl'.
  }
  eapply equiv_hasvalue; eauto.
  apply Htau.
  exists m.
  apply Hfg in Hf'.
  eenough (r = _) as -> by eauto.
  eapply hasvalue_det; eauto.
  apply eval_hasvalue.
Qed.

Lemma tree_fun_cont_self_modulus [A Q R : Type] `{D : EqDec R}
  (F : (Q -> part A) -> part R)
  (tau : list A -> part ( @result Q R ))
  (Htau : tree_fun_for tau F)
  : modulus (tree_fun_cont_mod_fn F tau Htau) (tree_fun_cont_mod_fn F tau Htau).
Proof.
  intros f (l & Hl).
  exists l.
  split; auto.
  intros _ g Hg.
  eapply equiv_hasvalue.
  1: eassumption.
  unfold tree_fun_cont_mod_fn in Hl |- *.
  rewrite bind_hasvalue in Hl |- *.
  destruct Hl as (Hf' & _ & Hl).
  edestruct (constructive_indefinite_ground_description_nat _ _ _)
  as (kf & Hkf).
  simpl in *.
  destruct (unembed kf) as (nf & mf).
  assert (Hf : eval_ext_tree tau f nf =! output (eval Hf'))
  by (rewrite seval_hasvalue; eauto).
  clear Hkf mf kf.
  assert (Hg' : ter (F g)). {
    destruct Hf' as (r & Hr).
    eexists.
    apply Htau.
    exists nf.
    destruct (tree_fun_trace_mod tau nf nil f ltac:(eexists; eauto))
    as (l' & Hl'' & Hl').
    assert (l' = l) as -> by (eapply hasvalue_det; eauto).
    specialize (Hl' ltac:(eexists;eauto) g Hg).
    apply Hl'.
    eauto.
  }
  destruct (proof_ter_spec (F g) Hg') as (Hg'' & H').
  exists Hg''.
  split; auto.
  clear H' Hg'.
  eset (cg := constructive_indefinite_ground_description_nat _).
  destruct cg as (kg & Hkg).
  clear cg.
  simpl.
  destruct (unembed kg) as (ng & mg).
  assert (Hng : eval_ext_tree tau g ng =! output (eval Hg''))
  by (rewrite seval_hasvalue; eauto).
  clear Hkg kg mg.
  assert (Hg' : eval_ext_tree tau g nf =! output (eval Hf')). {
    destruct (tree_fun_trace_mod _ nf _ _ ltac:(eexists; eauto))
    as (l' & Hl'' & Hl').
    assert (l' = l) by (eapply hasvalue_det; eauto); subst.
    specialize (Hl' ltac:(eexists; eauto) g Hg).
    simpl in Hl'.
    apply Hl'.
    assumption.
  }
  assert (E : eval Hf' = eval Hg'')
  by (eapply eval_ext_tree_aux_det; eauto).
  destruct E.
  eapply tree_fun_trace_aux_det.
  - eapply Hg'.
  - eapply Hng.
  - eapply tree_fun_trace_aux_mod.
    1: eapply Hf.
    all: eauto.
Qed.

Lemma tree_fun_cont_to_self_modulus_cont [Q A R : Type]
  `{D : EqDec R} (F : (Q -> part A) -> part R) :
  tree_fun_cont F -> self_modulus_cont F.
Proof.
  intros (tau & Htau).
  eexists.
  split.
  - unshelve eapply tree_fun_cont_modulus; eauto.
  - eapply tree_fun_cont_self_modulus; eauto.
Qed.

(* TODO : limite entre deux trucs différents ; est-ce que ça devrait être
dans un autre fichier ? *)


Definition neighborhood_fun_to_tree_fun [A R : Type]
  (gamma : neighborhood_fun A R) : @ext_tree nat A R :=
  fun l => bind (gamma l) (fun o => ret (match o with
  | Some r => output r
  | None => ask (length l)
  end)).

Lemma neighborhood_to_tree_trace [A R : Type] (gamma : neighborhood_fun A R)
  (f : nat -> part A) :
    forall (n : nat), ter (f n) ->
      (forall (k : nat), k <= n -> bind (fold_part_list (tabulate f 0 k)) gamma =! None)
      -> tree_fun_trace (neighborhood_fun_to_tree_fun gamma) f n =! tabulate id 0 n.
Proof.
  intros n Hn' Hn.
  induction n as [|n IHn].
  1: psimpl.
  rewrite tree_fun_trace_succ.
  assert (ter (f n)) as Hn''. {
    specialize (Hn (S n) ltac:(lia)).
    rewrite tabulate_succ', fold_part_list_app in Hn; try lia.
    cbn in Hn.
    psimpl.
    eexists; eauto.
  }
  specialize (IHn Hn'' (fun k Hk => Hn k ltac:(lia))).
  specialize (Hn n ltac:(lia)).
  simpl_assms.
  simpl in *.
  assert (n = length x) by now apply fold_part_list_tabulate_length in H.
  subst.
  destruct Hn''.
  psimpl.
  - unfold tabulate.
    rewrite tabulate_map.
    eauto.
  - rewrite <- tabulate_succ'; try lia.
    psimpl.
Qed.

Lemma neighborhood_fun_to_tree_fun_spec [A R : Type] :
  forall (gamma : neighborhood_fun A R) (f : nat -> part A) (n : nat),
    equiv (eval_ext_tree (neighborhood_fun_to_tree_fun gamma) f n)
      (mu_opt (fun k =>
        if k <=? n then
          bind (fold_part_list (tabulate f 0 k))
            (fun l => bind (neighborhood_fun_to_tree_fun gamma l) (fun res => ret (match res with
            | output r => Some (output r)
            | ask _ => None
            end)))
        else
          ret (Some (ask n))
      )).
Proof.
  intros gamma f n.
  induction n as [|n IHn].
  - cbn.
    unfold neighborhood_fun_to_tree_fun.
    intros res.
    split.
    + intros (o & Ho & Hres)%bind_hasvalue.
      eapply mu_opt_hasvalue.
      destruct o as [q|];
      [exists 0 | exists 1]; (split; [ psimpl | intros k Hk; try lia]).
      inversion Hk; subst.
      2: lia.
      rewrite leb_correct; eauto.
      cbn.
      psimpl.
    + intros (n & Hn & Hn')%mu_opt_hasvalue.
      destruct n as [|n]. 1: {
        cbn in Hn.
        psimpl.
        destruct x1; cbn;
        inversion H4; subst; psimpl.
      }
      specialize (Hn' 0 ltac:(lia)).
      cbn in Hn, Hn'.
      psimpl.
      destruct x1; psimpl.
  - rewrite eval_ext_tree_succ, IHn.
    clear IHn.
    intros res.
    split.
    + intros (res' & (m & Hm & Hm')%mu_opt_hasvalue & Hres')%bind_hasvalue.
      destruct (Nat.le_gt_cases m n) as [Hmn | Hmn].
      * rewrite leb_correct in Hm; try lia.
        apply mu_opt_hasvalue; exists m.
        rewrite leb_correct; try lia.
        unfold neighborhood_fun_to_tree_fun in *.
        setoid_rewrite equiv_bind_mult in Hm.
        setoid_rewrite equiv_bind_ret in Hm.
        simpl_assms.
        destruct x0; inversion H3; subst.
        simpl_assms.
        clear H3 H1.
        split; eauto.
        1: try now psimpl.
        intros k Hk.
        specialize (Hm' k Hk).
        rewrite leb_correct in Hm' |- *; try lia.
        setoid_rewrite equiv_bind_mult in Hm'.
        setoid_rewrite equiv_bind_ret in Hm'.
        setoid_rewrite equiv_bind_mult.
        setoid_rewrite equiv_bind_ret.
        psimpl.
      * assert (m = S n); subst. {
          inversion Hmn; subst; eauto.
          specialize (Hm' m0 ltac:(lia)).
          rewrite leb_correct_conv in Hm'; eauto.
          psimpl.
        }
        rewrite leb_correct_conv in Hm; try lia.
        simpl_assms.
        unfold neighborhood_fun_to_tree_fun in Hm', H2 |- *.
        simpl_assms.
        clear H4 Hmn.
        assert (x0 = tabulate id 0 n). {
          eapply hasvalue_det; eauto.
          apply neighborhood_to_tree_trace.
          1: eexists; eauto.
          intros k Hk.
          specialize (Hm' k ltac:(lia)).
          rewrite leb_correct in Hm'; eauto.
          psimpl.
          destruct x5; now inversion H8.
        }
        subst.
        unfold tabulate in H1.
        rewrite tabulate_map in H1.
        simpl in H1.
        destruct x2 as [r|]; apply mu_opt_hasvalue.
        --exists (S n).
          rewrite leb_correct; try lia.
          rewrite tabulate_succ', fold_part_list_app; try lia.
          cbn.
          split.
          1: psimpl.
          intros k Hk.
          rewrite leb_correct; try lia.
          specialize (Hm' k Hk).
          rewrite leb_correct in Hm'; try lia.
          psimpl.
        --exists (S (S n)).
          rewrite leb_correct_conv; try lia.
          split. 1: {
            enough (n = length x1) by (subst; psimpl).
            now apply fold_part_list_tabulate_length in H1.
          }
          intros k Hk.
          rewrite leb_correct; try lia.
          inversion Hk; subst.
          ++rewrite tabulate_succ', fold_part_list_app; try lia.
            cbn.
            psimpl.
          ++specialize (Hm' k H4).
            rewrite leb_correct in Hm'; try lia.
            psimpl.
    + intros (m & Hm & Hm')%mu_opt_hasvalue.
      destruct (Nat.le_gt_cases m (S n)) as [Hmn | Hmn].
      * rewrite leb_correct in Hm; eauto.
        simpl_assms.
        destruct x0; inversion H3; subst.
        clear H3.
        inversion Hmn; subst.
        --rewrite tabulate_succ', fold_part_list_app in H; try lia.
          cbn in H.
          simpl_assms.
          clear H3.
          cbn in H0.
          eenough (mu_opt _ =! ask n). {
            rewrite bind_hasvalue_given; eauto; simpl.
            rewrite bind_hasvalue_given; eauto; simpl.
            rewrite bind_hasvalue_given with (a := tabulate id 0 n).
            2: {
              apply neighborhood_to_tree_trace.
              1: eexists; eauto.
              intros k Hk.
              specialize (Hm' k ltac:(lia)).
              rewrite leb_correct in Hm'; try lia.
              unfold neighborhood_fun_to_tree_fun in Hm'.
              psimpl.
              destruct x3; now inversion H8.
            }
            unfold tabulate.
            rewrite tabulate_map.
            rewrite bind_hasvalue_given; eauto.
          }
          apply mu_opt_hasvalue.
          exists (S n).
          rewrite leb_correct_conv; try lia.
          split.
          1: psimpl.
          intros k Hk.
          specialize (Hm' k Hk).
          rewrite leb_correct in Hm' |- *; try lia.
          eauto.
        --eenough (mu_opt _ =! output r)
          by (rewrite bind_hasvalue_given; eauto; psimpl).
          apply mu_opt_hasvalue.
          exists m.
          rewrite leb_correct; eauto.
          unfold neighborhood_fun_to_tree_fun in H0.
          simpl_assms.
          destruct x0; inversion H4; subst.
          clear H4.
          split.
          1:psimpl.
          intros k Hk.
          rewrite leb_correct; try lia.
          specialize (Hm' k Hk).
          rewrite leb_correct in Hm'; try lia.
          psimpl.
      * rewrite leb_correct_conv in Hm; eauto.
        simpl_assms.
        specialize (Hm' (S n) ltac:(lia)) as Hn.
        rewrite Nat.leb_refl in Hn.
        rewrite tabulate_succ', fold_part_list_app in Hn; try lia.
        cbn in Hn.
        simpl_assms.
        destruct x0; inversion H6; subst.
        clear H6 H4.
        simpl in H0.
        assert (tree_fun_trace (neighborhood_fun_to_tree_fun gamma) f n =! tabulate id 0 n). {
          apply neighborhood_to_tree_trace.
          1: eexists; eauto.
          intros k Hk.
          specialize (Hm' k ltac:(lia)).
          rewrite leb_correct in Hm'; try lia.
          unfold neighborhood_fun_to_tree_fun in Hm'.
          psimpl.
          destruct x2; now inversion H7.
        }
        eenough (mu_opt _ =! ask n). {
          unfold neighborhood_fun_to_tree_fun in H0.
          rewrite bind_hasvalue_given; eauto; psimpl.
          - unfold tabulate; rewrite tabulate_map; eauto.
          - destruct x; inversion H6; subst.
            enough (n = length x1) by (subst; psimpl).
            now apply fold_part_list_tabulate_length in H.
        }
        eapply mu_opt_hasvalue.
        exists (S n).
        rewrite leb_correct_conv; try lia.
        split.
        1: psimpl.
        intros k Hk.
        specialize (Hm' k ltac: (lia)).
        rewrite leb_correct in Hm' |- *; try lia.
        eauto.
Qed.

Lemma associated_neighborhood_fun_to_associated_tree_fun [A R : Type] (F : (nat -> part A) -> part R)
  (gamma : neighborhood_fun A R) : neighborhood_for gamma F ->
      tree_fun_for (neighborhood_fun_to_tree_fun gamma) F.
Proof.
  intros Hgamma.
  intros f r.
  setoid_rewrite neighborhood_fun_to_tree_fun_spec.
  split; intros Hr.
  - apply Hgamma in Hr as (n & Hn & Hn')%mu_opt_hasvalue.
    exists n.
    eapply mu_opt_hasvalue.
    exists n.
    rewrite leb_correct; try eauto.
    psimpl.
    intros k Hk; rewrite leb_correct; try lia.
    specialize (Hn' k Hk).
    psimpl.
  - apply Hgamma.
    destruct Hr as (n & (m & Hm & Hm')%mu_opt_hasvalue).
    destruct (Nat.le_gt_cases m n) as [Hmn | Hmn].
    + rewrite leb_correct in Hm; eauto.
      simpl_assms.
      destruct x0; inversion H3.
      unfold neighborhood_fun_to_tree_fun in H0; simpl_assms.
      destruct x0; inversion H5; subst.
      apply mu_opt_hasvalue.
      exists m; split; try now psimpl.
      intros k Hk.
      specialize (Hm' k ltac:(lia)).
      rewrite leb_correct in Hm'; try lia.
      unfold neighborhood_fun_to_tree_fun in Hm'.
      setoid_rewrite equiv_bind_mult in Hm'.
      setoid_rewrite equiv_bind_ret in Hm'.
      psimpl.
      destruct x1; psimpl.
    + rewrite leb_correct_conv in Hm; eauto.
      psimpl.
Qed.

Lemma neighborhood_cont_to_tree_fun_cont [A R : Type] :
  forall (F : (nat -> part A) -> part R),
    neighborhood_cont F -> tree_fun_cont F.
Proof.
  intros F (gamma & Hgamma).
  eexists.
  eapply associated_neighborhood_fun_to_associated_tree_fun.
  eauto.
Qed.

Lemma tree_fun_cont_not_neighborhood_cont :
  exists (F : (nat -> part unit) -> part unit),
    tree_fun_cont F /\ ~ neighborhood_cont F.
Proof.
  set (F (f : nat -> part unit) := f 1).
  exists F.
  split.
  - set (tau (l : list unit) := match l with
    | nil => ret (ask 1)
    | _ => ret (output ())
    end).
    exists tau.
    intros f ().
    split.
    + intros Hf.
      exists 1.
      cbn.
      psimpl.
    + intros (n & Hf).
      unfold F.
      destruct n; cbn in *; psimpl.
      destruct x; eauto.
  - intros (gamma & Hgamma).
    set (f n := match n with
    | 0 => undef
    | _ => ret ()
    end).
    destruct (Hgamma f ()) as (Hgammaf & _).
    specialize (Hgammaf ltac:(psimpl)) as (n & Hn & Hn')%mu_opt_hasvalue.
    destruct n.
    + cbn in Hn.
      rewrite equiv_bind_ret in Hn.
      destruct (Hgamma (fun _ => undef) ()) as (_ & Hgammaundef).
      assert (F (fun _ => undef) =! ()). {
        apply Hgammaundef.
        apply mu_opt_hasvalue; exists 0; split; psimpl.
      }
      unfold F in *.
      psimpl.
    + apply bind_hasvalue in Hn as (l & Hl & Hl').
      rewrite tabulate_succ, fold_part_list_cons, bind_hasvalue in Hl; try lia.
      destruct Hl as ([] & Hf & Hf').
      unfold f in Hf.
      psimpl.
Qed.

Definition tree_fun_cont_strong_modulus_fn' [Q A R : Type] (tau : @ext_tree Q A R) : (Q -> part A) -> part (list Q) :=
  fun f => mu_opt (fun n => bind (eval_ext_tree tau f n) (fun o =>
    match o with
    | ask _ => ret None
    | output r => bind (tree_fun_trace tau f n) (fun l => ret (Some l))
    end
  )).

Lemma eval_ext_tree_aux_prev_ask [Q A R : Type] (tau : @ext_tree Q A R) :
  forall (f : Q -> part A) (n : nat) (l : list A),
    (exists (q : Q), eval_ext_tree_aux tau f n l =! ask q) ->
      forall (k : nat), k <= n -> (exists (q : Q), eval_ext_tree_aux tau f k l =! ask q).
Proof.
  intros f n l Hn k Hk.
  induction Hk as [| n Hk IHk]; eauto.
  apply IHk.
  clear k Hk IHk.
  destruct Hn as (q & Hq).
  assert (ter (eval_ext_tree_aux tau f n l)) as ([q' | r] & Hr)
  by (eapply eval_ext_tree_aux_succ_ter; eexists; eauto); eauto.
  enough (ask q = output r) as [=].
  eapply hasvalue_det; eauto.
  eapply eval_ext_tree_aux_succ_result; eauto.
Qed.

Lemma eval_ext_tree_aux_equiv_trace [Q A R : Type] (tau : @ext_tree Q A R) :
  forall (f : Q -> part A) (n : nat) (l' : list A) (l : list Q) (o : result),
    eval_ext_tree_aux tau f n l' =! o ->
    tree_fun_trace_aux tau f n l' =! l -> forall (g : Q -> part A),
    Forall (fun q => equiv (f q) (g q)) l ->
    eval_ext_tree_aux tau g n l' =! o.
Proof.
  intros f n.
  induction n as [|n IHn]; intros l' l o Ho Hl' g Hg.
  1: now cbn in Ho |- *.
  cbn in *.
  psimpl.
  assert (x0 = x) by (eapply hasvalue_det; eauto); subst.
  destruct x; simpl_assms.
  assert (x1 = x) by (eapply hasvalue_det; eauto); subst.
  inversion Hg; subst.
  psimpl.
  now apply H7.
Qed.

Lemma tree_fun_trace_aux_equiv_trace [Q A R : Type] (tau : @ext_tree Q A R) :
  forall (f : Q -> part A) (n : nat) (l' : list A) (l : list Q) (o : result),
    eval_ext_tree_aux tau f n l' =! o ->
    tree_fun_trace_aux tau f n l' =! l -> forall (g : Q -> part A),
    Forall (fun q => equiv (f q) (g q)) l ->
    tree_fun_trace_aux tau g n l' =! l.
Proof.
  intros f n.
  induction n as [|n IHn]; intros l' l o Ho Hl' g Hg.
  1: now cbn in Ho |- *.
  cbn in *.
  psimpl.
  assert (x0 = x) by (eapply hasvalue_det; eauto); subst.
  destruct x; simpl_assms.
  assert (x1 = x) by (eapply hasvalue_det; eauto); subst.
  inversion Hg; subst.
  psimpl.
  now apply H7.
Qed.

Lemma tree_fun_trace_aux_restrict' [Q A R : Type] (tau : @ext_tree Q A R) (f : Q -> part A) :
  forall (n : nat) (l' : list A) (l : list Q),
    tree_fun_trace_aux tau f (S n) l' =! l ->
      exists (l1 l2 : list Q), tree_fun_trace_aux tau f n l' =! l2 /\ l = l2 ++ l1.
Proof.
  intros n.
  induction n as [|n IHn]; intros l' l Hn.
  1: exists l, nil; split; psimpl.
  change (tree_fun_trace_aux tau f (S (S n)) l') with
  (bind (tau l') (fun o => match o with
  | output _ => ret nil
  | ask q => let p := bind (f q) (fun a => tree_fun_trace_aux tau f (S n) (a :: l'))
      in bind p (fun l => ret (q :: l))
  end)) in Hn.
  simpl in |- *.
  apply bind_hasvalue in Hn as (res & Hq & Hn).
  setoid_rewrite bind_hasvalue_given; eauto.
  destruct res as [q | r].
  2: exists nil, l; split; eauto using app_nil_r.
  apply equiv_bind_mult in Hn.
  setoid_rewrite equiv_bind_mult.
  apply bind_hasvalue in Hn as (a & Ha & (l1' & Hn & Hl%ret_hasvalue_iff)%bind_hasvalue); subst.
  setoid_rewrite bind_hasvalue_given; eauto.
  apply IHn in Hn as (l1 & l2 & Hn & Hl1'); subst.
  eexists; psimpl.
Qed.

Lemma tree_fun_trace_aux_restrict [Q A R : Type] (tau : @ext_tree Q A R) (f : Q -> part A) :
  forall (n : nat) (l' : list A) (l : list Q),
    tree_fun_trace_aux tau f n l' =! l ->
      forall (k : nat), k <= n -> exists (l1 l2 : list Q),
        tree_fun_trace_aux tau f k l' =! l2 /\ l = l2 ++ l1.
Proof.
  intros n l' l Hn k Hk.
  revert l' l Hn.
  induction Hk as [| n Hk IHk]; intros l' l Hn.
  1: exists nil, l; eauto using app_nil_r.
  apply tree_fun_trace_aux_restrict' in Hn as (l0 & l1 & Hl1 & ?); subst.
  apply IHk in Hl1 as (l2 & l3 & Hl3 & ?); subst.
  rewrite <- app_assoc.
  eexists _; eauto.
Qed.

Lemma tree_fun_cont_strong_modulus' [Q A R : Type] (F : (Q -> part A) -> part R)
  (tau : @ext_tree Q A R) : tree_fun_for tau F ->
    strong_modulus F (tree_fun_cont_strong_modulus_fn' tau).
Proof.
  intros Htau f r Hr.
  specialize (Htau f r) as Htauf.
  pose (Hr' := Hr).
  apply Htauf in Hr' as (n & Hn).
  assert (ter (tree_fun_trace tau f n)) as (l & Hl)
  by (apply tree_fun_trace_ter; eexists; eauto).
  exists l.
  repeat split.
  - unfold tree_fun_cont_strong_modulus_fn'.
    generalize dependent l.
    induction n as [|n IHn]; intros l Hl.
    + apply mu_opt_hasvalue.
      exists 0.
      split.
      2: intros k Hk; inversion Hk.
      psimpl.
      cbn in Hl.
      apply ret_hasvalue_inv in Hl;
      rewrite Hl; psimpl.
    + assert (ter (eval_ext_tree tau f n)) as [[q | r'] Hn']
      by (apply eval_ext_tree_aux_succ_ter; eexists; eauto).
      * clear IHn.
        eapply mu_opt_hasvalue.
        exists (S n); split.
        1: repeat (rewrite bind_hasvalue_given; eauto; simpl); psimpl.
        intros k Hk.
        assert (exists (q : Q), eval_ext_tree tau f k =! ask q) as (q' & Hq')
        by (eapply eval_ext_tree_aux_prev_ask; eauto; try lia).
        psimpl.
      * assert (r' = r) by (eapply eval_ext_tree_aux_det; eauto); subst.
        assert (tree_fun_trace tau f n =! l) by
        (eapply tree_fun_trace_aux_det with (n := S n); eauto).
        eauto.
  - generalize dependent l.
    induction n as [|n IHn]; intros l Hl.
    1:cbn in Hl; apply ret_hasvalue_iff in Hl; subst; auto.
    rewrite <- Forall_map, <- fold_part_list_ter.
    eapply tree_fun_trace_ter_fold; try eexists; eauto.
  - intros g Hg.
    apply (Htau g r).
    exists n.
    eapply eval_ext_tree_aux_equiv_trace; eauto.
Qed.

Lemma tree_fun_cont_strong_self_modulus [Q A R : Type] (F : (Q -> part A) -> part R)
  (tau : @ext_tree Q A R) : tree_fun_for tau F ->
  strong_modulus (tree_fun_cont_strong_modulus_fn' tau) (tree_fun_cont_strong_modulus_fn' tau).
Proof.
  intros Htau f l Hl.
  eexists; split; eauto.
  split.
  - unfold tree_fun_cont_strong_modulus_fn' in Hl.
    eapply mu_opt_hasvalue in Hl as (n & Hl & _).
    psimpl.
    destruct x; psimpl.
    rewrite <- Forall_map, <- fold_part_list_ter.
    eapply tree_fun_trace_ter_fold; try eexists; eauto.
  - intros g Hg.
    unfold tree_fun_cont_strong_modulus_fn' in Hl |- *.
    apply mu_opt_hasvalue in Hl as (n & Hl & Hl').
    apply mu_opt_hasvalue; exists n.
    simpl_assms.
    destruct x; simpl_assms.
    rename H into Hr, H0 into Hn.
    assert (eval_ext_tree tau g n =! output r)
    by (eauto using eval_ext_tree_aux_equiv_trace).
    assert (tree_fun_trace tau g n =! l)
    by (eauto using tree_fun_trace_aux_equiv_trace).
    repeat (rewrite bind_hasvalue_given; eauto; simpl).
    split; try now psimpl.
    intros k Hkn.
    specialize (Hl' k Hkn).
    simpl_assms.
    destruct x; simpl_assms.
    rename H into Hr', H0 into Hn', H1 into Hq.
    eapply tree_fun_trace_aux_restrict with (k := k) in Hn as (l1 & l2 & Hk & ?); try lia; subst.
    apply Forall_app_l in Hg.
    assert (eval_ext_tree tau g k =! ask q)
    by (eauto using eval_ext_tree_aux_equiv_trace).
    assert (tree_fun_trace tau g k =! l2)
    by (eauto using tree_fun_trace_aux_equiv_trace).
    repeat (rewrite bind_hasvalue_given; eauto; simpl).
    psimpl.
Qed.

(* TODO : à garder ? *)
Lemma tree_fun_for_good_def [Q A R : Type] :
  forall (F : (Q -> part A) -> part R) (tau : @ext_tree Q A R),
    tree_fun_for tau F <->
      (forall f : Q A, (ter (A:=R) (F f) -> exists n : nat,
        equiv (A:=result) (eval_ext_tree tau f n)
        (bind (B:=result) (F f) (fun r : R => ret (output r)))) /\
        ((exists (n : nat) (r : R), eval_ext_tree tau f n =! output r) ->
          ter (A:=R) (F f))).
Proof.
  intros F tau.
  split.
  - intros Htau f.
    specialize (Htau f).
    split.
    + intros (r & Hr).
      pose proof (Hr' := Hr).
      apply (Htau r) in Hr' as (n & Hn).
      eexists.
      eapply equiv_hasvalue; psimpl.
    + intros (n & r & Hr).
      eexists.
      eapply Htau.
      eauto.
  - intros Htau f r.
    specialize (Htau f) as (Htau1 & Htau2).
    split.
    + intros Hr.
      specialize (Htau1 ltac:(eexists; eauto)) as (n & Hn).
      eexists; apply Hn.
      psimpl.
    + intros (n & Hn).
      specialize (Htau2 ltac:(exists n, r; apply Hn)) as (r' & Hr).
      specialize (Htau1 ltac:(eexists; eauto)) as (n' & Hn').
      assert (eval_ext_tree tau f n' =! output r') by (apply Hn'; psimpl).
      enough (r = r') by now subst.
      eapply eval_ext_tree_aux_det; eauto.
Qed.

Theorem tree_fun_cont_to_strong_self_modulus_cont [Q A R : Type] :
  forall (F : (Q -> part A) -> part R),
    tree_fun_cont F -> strong_self_modulus_cont F.
Proof.
  intros F (tau & Htau).
  eexists.
  split.
  - eapply tree_fun_cont_strong_modulus'; eauto.
  - eapply tree_fun_cont_strong_self_modulus; eauto.
Qed.

(* TODO : je crois que ça marche aussi avec (bool -> part unit) -> part unit *)
Theorem self_mod_cont_not_tree_fun_cont :
  exists (F : (bool -> part bool) -> part bool),
    self_modulus_cont F /\ ~ tree_fun_cont F.
Proof.
  set (F f := por (f true) (f false)).
  exists F.
  split.
  - set (M (_ : bool -> part bool) := ret [true ; false]).
    exists M.
    split.
    + intros f Hf.
      eexists; split; psimpl.
      intros _ g Hg.
      inversion Hg.
      inversion H2.
      clear H4 H3 l0 x0 H0 H H2 l x.
      unfold F.
      rewrite H1, H5.
      reflexivity.
    + intros f Hf.
      exists [true ; false ]; split; psimpl.
  - intros (tau & Htau).
    set (f1 (b : bool) := if b then ret true else undef).
    set (f2 (b : bool) := if b then undef else ret true).
    set (f3 (b : bool) := @undef _ bool).
    assert (H1 : F f1 =! true). {
      unfold F.
      simpl.
      rewrite por_true.
      left. psimpl.
    }
    assert (H2 : F f2 =! true). {
      unfold F.
      simpl.
      rewrite por_true.
      right. psimpl.
    }
    assert (~ter (F f3)). {
      unfold F.
      simpl.
      intros ([] & Hb).
      - rewrite por_true in Hb.
        destruct Hb as [|];
        eapply undef_hasvalue; eassumption.
      - rewrite por_false in Hb;
        eapply undef_hasvalue; destruct Hb; eassumption.
    }
    pose proof H1 as (n & H1')%Htau.
    pose proof H2 as (m & H2')%Htau.

    assert (Hc : forall b, tau nil =! output b -> False). {
      intros b Hb.
      apply H.
      exists b.
      apply Htau.
      exists 0.
      cbn.
      auto.
    }

    destruct n as [|n].
    1: eapply Hc, H1'; psimpl.
    destruct m as [|m].
    1: eapply Hc, H2'; psimpl.
    cbn in *.
    apply bind_hasvalue in H1' as (r & Hr & H1').
    destruct r; try now (eapply Hc; eauto).
    rewrite bind_hasvalue_given in H2'; eauto.
    cbn in *.
    psimpl.
    unfold f1, f2 in *.
    destruct b; psimpl.
Qed.

End PartialZoo.