From Stdlib Require Import Lia Arith Program Setoids.Setoid Classes.Morphisms.
From Stdlib Require Lists.List.
Import List.ListNotations.
Open Scope list_scope.

Require Import continuity_zoo_Prop (ext_tree, result(..)).
Require Import Partial.TreeFunction Partial.Delay Partial.Common partial.
Import PartialTactics DelayPartialImplem.

Set Bullet Behavior "Strict Subproofs".

Unset Elimination Schemes.
Inductive dialogue_tree (Q A R : Type) : Type :=
| Output : R -> dialogue_tree Q A R
| Ask : Q -> (A -> delay (dialogue_tree Q A R)) -> dialogue_tree Q A R.
Set Elimination Schemes.

Arguments Ask [_ _ _].
Arguments Output [_ _ _].

(* TODO: sort poly ? *)
Fixpoint dialogue_tree_rect [Q A R : Type] (P : dialogue_tree Q A R -> Type)
  (HOutput : forall (r : R), P (Output r))
  (HAsk : forall (q : Q) (k : A -> delay (dialogue_tree Q A R)),
    (forall (a : A) (b : dialogue_tree Q A R), hasvalue (k a) b -> P b)
    -> P (Ask q k))
  (b : dialogue_tree Q A R) : P b.
Proof.
  destruct b as [r | q k].
  - apply HOutput.
  - apply HAsk.
    intros a b Hab.
    induction Hab as [b | d b Hd IHd].
    + apply dialogue_tree_rect; auto.
    + apply IHd.
Defined.

Definition dialogue_tree_ind [Q A R : Type] (P : dialogue_tree Q A R -> Prop)
  (HOutput : forall (r : R), P (Output r))
  (HAsk : forall (q : Q) (k : A -> delay (dialogue_tree Q A R)),
    (forall (a : A) (b : dialogue_tree Q A R), hasvalue (k a) b -> P b)
    -> P (Ask q k)) :
  forall (b : dialogue_tree Q A R), P b
:= dialogue_tree_rect P HOutput HAsk.

Lemma dialogue_tree_rect_output (Q A R : Type) (P : dialogue_tree Q A R -> Type)
  (HAsk : forall (q : Q) (k : A -> delay (dialogue_tree Q A R)),
    (forall (a : A) (b : dialogue_tree Q A R), hasvalue (k a) b -> P b)
    -> P (Ask q k))
  (HOutput : forall (r : R), P (Output r)) :
  forall (r : R), dialogue_tree_rect P HOutput HAsk (Output r) = HOutput r.
Proof.
  intros r.
  reflexivity.
Qed.

Lemma dialogue_tree_rect_ask (Q A R : Type) (P : dialogue_tree Q A R -> Type)
  (HAsk : forall (q : Q) (k : A -> delay (dialogue_tree Q A R)),
    (forall (a : A) (b : dialogue_tree Q A R), hasvalue (k a) b -> P b)
    -> P (Ask q k))
  (HOutput : forall (r : R), P (Output r)) :
  forall (q : Q) (k : A -> delay (dialogue_tree Q A R)),
    dialogue_tree_rect P HOutput HAsk (Ask q k) =
      HAsk q k (fun a b Hab => dialogue_tree_rect P HOutput HAsk b).
Proof.
  intros q k.
  simpl.
  f_equal.
  (* TODO : faire sans *)
  extensionality a.
  extensionality b.
  extensionality Hab.
  induction Hab as [ b | d b Hb IHb ].
  - cbn.
    reflexivity.
  - cbn in *.
    apply IHb.
Qed.

Definition eval_dialogue_tree [Q A R : Type] (d : dialogue_tree Q A R) (f : Q -> delay A) : delay R.
Proof.
  induction d as [ r | q k Hk ].
  - exact (Now r).
  - exact (bind (proof_ter (f q)) (fun Ha =>
        bind (proof_ter (k (eval Ha))) (fun Hb =>
          Hk (eval Ha) (eval Hb) (eval_hasvalue Hb)))).
Defined.

Lemma eval_dialogue_tree_output [Q A R : Type] (f : Q -> delay A) :
  forall (r : R), hasvalue (eval_dialogue_tree (Output r) f) r.
Proof.
  intros r.
  cbn.
  constructor.
Qed.

Lemma eval_dialogue_tree_ask [Q A R : Type] (q : Q) (k : A -> delay (dialogue_tree Q A R))
  (f : Q -> delay A) :
  equiv (eval_dialogue_tree (Ask q k) f) (bind (f q) (fun a =>
    bind (k a) (fun d => eval_dialogue_tree d f))).
Proof.
  intros r.
  split; intros Hr.
  - unfold eval_dialogue_tree in Hr.
    rewrite dialogue_tree_rect_ask in Hr.
    apply bind_hasvalue in Hr as (Hq & _ & (Ha & _ & Hr)%bind_hasvalue).
    rewrite bind_hasvalue_given with (a := eval Hq).
    2: apply eval_hasvalue.
    rewrite bind_hasvalue_given with (a := eval Ha).
    2: apply eval_hasvalue.
    remember (eval Ha) as d eqn:E.
    induction d as [ r' | q' k' IHk' ].
    + cbn in Hr |- *.
      assumption.
    + unfold eval_dialogue_tree.
      assumption.
  - apply bind_hasvalue in Hr as (a & Ha & (d & Hd & Hr)%bind_hasvalue).
    unfold eval_dialogue_tree.
    rewrite dialogue_tree_rect_ask.
    assert (ter (proof_ter (f q))) as (Hater & Hater')
    by (apply proof_ter_spec; eexists; eassumption).
    rewrite bind_hasvalue_given.
    2: eassumption.
    assert (Ea : eval Hater = a) by
    (eapply partial.hasvalue_det; [ apply eval_hasvalue | assumption ]).
    rewrite Ea.
    assert (ter (proof_ter (k a))) as (Hdter & Hdter')
    by (apply proof_ter_spec; eexists; try eassumption).
    rewrite bind_hasvalue_given.
    2: eassumption.
    assert (Ed : eval Hdter = d) by
    (eapply partial.hasvalue_det; [ apply eval_hasvalue | assumption ]).
    rewrite Ed.
    induction d as [ r' | q' k' IHk' ].
    + cbn in Hr |- *.
      assumption.
    + unfold eval_dialogue_tree in Hr.
      apply Hr.
Qed.

Definition dialogue_tree_for [Q A R : Type] (tau : dialogue_tree Q A R) (F : (Q -> delay A) -> delay R) : Prop :=
  forall (f : Q -> delay A), equiv (eval_dialogue_tree tau f) (F f).

Definition dialogue_tree_cont [Q A R : Type] (F : (Q -> delay A) -> delay R) : Prop :=
  exists (tau : dialogue_tree Q A R), dialogue_tree_for tau F.

(* TODO : adapter pour prendre un tree_fun_for *)
(* TODO : est-ce que c'est encore équivalent à la version de neighborhood tree ? *)
#[projections(primitive)]
Record sigma1_tree (Q A : Type) := mkSigma1Tree {
  ptree_branch : (Q -> delay A) -> Prop;
  ptree_path : list (Q * A) -> Prop;
  F : (Q -> delay A) -> delay unit;
  tau : @ext_tree _ Q A unit;
  Htau : strict_well_founded_tree_fun_for tau F;
  Hptree_branch : reflectspfn ptree_branch F;
  Hptree_path : reflectspfn ptree_path (fun l => bind (tau (snd (List.split l))) (fun _ => ret ()));
}.

Arguments ptree_branch [_ _].
Arguments ptree_path [_ _].
Arguments F [_ _].
Arguments tau [_ _].

Lemma tree_fun_for_squash [Q A R : Type] :
  forall (F : (Q -> delay A) -> delay R) (tau : @ext_tree _ Q A R),
    tree_fun_for tau F -> tree_fun_for (fun l => bind (tau l) (fun o => ret match o with
    | output _ => output ()
    | ask q => ask q
    end)) (fun f => bind (F f) (fun _ => ret ())).
Proof.
  intros F tau Htau.
  intros f.
  enough (Heq : forall (n : nat) (l : list A), equiv (bind (eval_ext_tree_aux tau f n l) (fun o => ret match o with
    | ask q => ask q
    | output _ => output ()
    end))
    (eval_ext_tree_aux (fun l => bind (tau l) (fun o => ret match o with
    | ask q => ask q
    | output _ => output ()
    end)) f n l)
  ). {
    intros ().
    setoid_rewrite <- Heq.
    split.
    - intros (r & (n & Hn)%Htau & _)%bind_hasvalue.
      exists n.
      rewrite bind_hasvalue_given; eauto.
      constructor.
    - intros (n & ([q | r] & Hr & Hn)%bind_hasvalue); cbn in Hn; try now inversion Hn.
      assert (F f =! r) by (apply Htau; eexists; eassumption).
      rewrite bind_hasvalue_given; eauto.
      constructor.
  }
  intros n.
  induction n as [|n IHn]; intros l.
  - cbn.
    reflexivity.
  - cbn.
    repeat change (bind ?a ?f) with (partial.bind a f).
    setoid_rewrite equiv_bind_mult.
    setoid_rewrite equiv_bind_ret.
    apply equiv_bind.
    1: reflexivity.
    intros r Hr _.
    destruct r as [q | r].
    2: rewrite equiv_bind_ret; reflexivity.
    rewrite equiv_bind_mult.
    setoid_rewrite IHn.
    reflexivity.
Qed.

Lemma well_founded_tree_fun_for_squash [Q A R : Type] :
  forall (F : (Q -> part A) -> part R) (tau : @ext_tree _ Q A R),
    well_founded_tree_fun_for tau F -> well_founded_tree_fun_for (fun l => bind (tau l) (fun o => ret match o with
    | output _ => output ()
    | ask q => ask q
    end)) (fun f => bind (F f) (fun _ => ret ())).
Proof.
  intros F tau ((Hwf1 & Hwf2) & Htau).
  split.
  2: apply tree_fun_for_squash; eauto.
  split.
  - intros l l' (_ & (res & Hres & _)%bind_hasvalue).
      assert (ter (tau l)) as (o & Ho) by (eapply Hwf1; eexists; eassumption).
      eexists.
      rewrite bind_hasvalue_given.
      2: eassumption.
      apply ret_hasvalue.
    - intros () l l' (res & Hres & Hres')%bind_hasvalue.
      change (hasvalue (ret ?a) ?a') with (partial.ret a =! a') in Hres'.
      apply ret_hasvalue_iff in Hres'.
      destruct res; inversion Hres'.
      eapply Hwf2 in Hres.
      rewrite bind_hasvalue_given.
      2: eassumption.
      apply ret_hasvalue.
Qed.

Lemma strict_well_founded_tree_fun_for_squash [Q A R : Type] :
  forall (F : (Q -> part A) -> part R) (tau : @ext_tree _ Q A R),
    strict_well_founded_tree_fun_for tau F -> strict_well_founded_tree_fun_for (fun l => bind (tau l) (fun o => ret match o with
    | output _ => output ()
    | ask q => ask q
    end)) (fun f => bind (F f) (fun _ => ret ())).
Proof.
  intros F tau (Hstrict & Htau).
  split.
  2: apply well_founded_tree_fun_for_squash; eauto.
  intros l.
  split.
  - intros (res' & (res & Hres & Hres')%bind_hasvalue).
    assert (exists (r : R) (l' : list A), (tau (l' ++ l)) =! output r)
    as (r & l' & Hr) by (eapply Hstrict; eexists; eassumption).
    exists (), l'.
    rewrite bind_hasvalue_given.
    2: eassumption.
    apply ret_hasvalue.
  - intros (() & l' & (res' & Hres' & Hr)%bind_hasvalue).
    destruct res'; cbn in Hr; inversion Hr; subst.
    enough (ter (tau l)) as (res & Hres)
    by (eexists; rewrite bind_hasvalue_given; [ apply ret_hasvalue | eassumption]).
    destruct Htau as ((Hwf1 & _) & _).
    eapply Hwf1.
    eexists; eauto.
Qed.

Definition termination_tree [Q A R : Type] (F : (Q -> delay A) -> delay R)
  (tau : ext_tree) :
    strict_well_founded_tree_fun_for tau F -> sigma1_tree Q A.
Proof.
  intros Htau.
  set (ptree_branch (f : Q -> delay A) := ter (Part := delay_monad) (F f)).
  set (ptree_path (l : list (Q * A)) := ter (Part := delay_monad) (tau (snd (List.split l)))).
  apply mkSigma1Tree with (ptree_branch := ptree_branch) (ptree_path := ptree_path)
    (F := fun f => bind (F f) (fun _ => ret ()))
    (tau := fun l => bind (tau l) (fun o => ret match o with
    | output _ => output ()
    | ask q => ask q
    end)).
  - eapply strict_well_founded_tree_fun_for_squash; assumption.
  - unfold reflectspfn, ptree_branch;
    intros f; split; intros (r & Hr); psimpl.
    + eexists; simpl_goal.
      apply ret_hasvalue.
    + apply bind_hasvalue in Hr as (a & Ha & _).
      eexists; eassumption.
  - unfold reflectspfn, ptree_branch;
    intros f; split; intros (r & Hr); psimpl.
    + eexists.
      rewrite bind_hasvalue_given.
      1: apply ret_hasvalue.
      rewrite bind_hasvalue_given.
      2: eassumption.
      apply ret_hasvalue.
    + change (bind (A := ?A) (B := ?B)) with (partial.bind (partiality := delay_monad) (A := A) (B := B)) in Hr.
      rewrite equiv_bind_mult in Hr.
      apply partial.bind_hasvalue in Hr as (r' & Hr' & _).
      eexists; eassumption.
Defined.

Definition pbarred [Q A : Type] (T : sigma1_tree Q A) (P : list (Q * A) -> Prop) : Prop :=
  forall (f : Q -> delay A), T.(ptree_branch) f ->
    exists (l : list (Q * A)), List.Forall (fun '(q, a) => (f q) =! a) l /\ P l.

(* TODO : export *)
Lemma split_app {A B : Type} (l l' : list (A * B)) :
  List.split (l ++ l') =
    let '(l1, l2) := List.split l in
    let '(l1', l2') := List.split l' in
    (l1 ++ l1', l2 ++ l2').
Proof.
  induction l as [| (a, b) l IHl ]; cbn.
  - destruct (List.split l'); reflexivity.
  - destruct (List.split (l ++ l')), (List.split l), (List.split l').
    inversion IHl; subst.
    reflexivity.
Qed.

(* TODO: export *)
Lemma cons_no_dup {A B : Type} (l : list (A * B)) (a : A) (b : B) :
  ~List.In a (fst (List.split l)) -> List.NoDup (fst (List.split l)) ->
  List.NoDup (fst (List.split ((a, b) :: l))).
Proof.
  intros Ha Hl.
  cbn.
  destruct (List.split l).
  cbn.
  constructor; auto.
Qed.

Lemma combine_app_same_length_r {A B : Type} (la1 la2 : list A) (lb1 lb2 : list B) :
  length la1 = length lb1 ->
    List.combine (la1 ++ la2) (lb1 ++ lb2) = List.combine la1 lb1 ++ List.combine la2 lb2.
Proof.
  intros E.
  remember (length lb1) as n eqn:F.
  symmetry in F.
  revert n la1 la2 lb1 lb2 E F.
  induction n as [|n IHn]; intros la1 la2 lb1 lb2 E F.
  - apply List.length_zero_iff_nil in E as ->.
    apply List.length_zero_iff_nil in F as ->.
    reflexivity.
  - destruct la1; cbn in E; inversion E.
    destruct lb1; cbn in F; inversion F.
    simpl.
    f_equal.
    apply IHn; assumption.
Qed.

(* TODO: export *)
Lemma rev_combine {A B : Type} (l1 : list A) (l2 : list B) :
  length l1 = length l2 ->
  List.rev (List.combine l1 l2) = List.combine (List.rev l1) (List.rev l2).
Proof.
  revert l2.
  induction l1 as [| a l1 IHl1 ] using List.rev_ind; intros l2 E; cbn.
  - reflexivity.
  - rewrite List.length_app, Nat.add_comm in E.
    rename l2 into l2'.
    assert ({ l2 : list B & { b : B | l2' = l2 ++ [b] }}) as (l2 & b & ->). {
      apply List.exists_last.
      1: intros ->; cbn in E; inversion E.
    }
    rewrite List.length_app, Nat.add_comm with (m := length [b]) in E.
    simpl in E.
    inversion E.
    rewrite combine_app_same_length_r; try assumption.
    rewrite 3! List.rev_app_distr.
    simpl.
    f_equal.
    apply IHl1.
    assumption.
Qed.

(* question que je me pose : je ne devrais pas, en plus, supposer que
  chaque q apparaît de façon unique dans l ? *)

Unset Elimination Schemes.
Inductive hereditarily_pbarred [Q A : Type] (T : sigma1_tree Q A) (P : list (Q * A) -> Prop)
  (l : list (Q * A)) (Hl : T.(ptree_path) l) (Hl' : List.NoDup (fst (List.split l))) : Prop :=
| Hpbar : P l \/ (~ P l /\
    forall (q : Q) (a : A) (Hal : T.(ptree_path) ((q, a) :: l))
      (Hq : ~ List.In q (fst (List.split l))),
        hereditarily_pbarred T P ((q, a) :: l) Hal (cons_no_dup _ _ _ Hq Hl')
  )
-> hereditarily_pbarred T P l Hl Hl'.
Set Elimination Schemes.

Fixpoint hereditarily_pbarred_ind [Q A : Type] (T : sigma1_tree Q A) (P : list (Q * A) -> Prop)
  (Qind : forall (l : list (Q * A)), T.(ptree_path) l -> (List.NoDup (fst (List.split l))) -> Prop) :
  (forall (l : list (Q * A)) (Hl : T.(ptree_path) l) (Hl' : List.NoDup (fst (List.split l))),
    (P l -> Qind l Hl Hl') * (
      ~ P l ->
      (forall (q : Q) (a : A) (Hal : T.(ptree_path) ((q, a) :: l))
        (Hq : ~ List.In q (fst (List.split l))),
          hereditarily_pbarred T P ((q, a) :: l) Hal (cons_no_dup _ _ _ Hq Hl')) ->
      (forall (q : Q) (a : A) (Hal : T.(ptree_path) ((q, a) :: l))
        (Hq : ~ List.In q (fst (List.split l))),
          Qind ((q, a) :: l) Hal (cons_no_dup _ _ _ Hq Hl')) ->
    Qind l Hl Hl')
  ) -> forall (l : list (Q * A)) (Hl : T.(ptree_path) l) (Hl' : List.NoDup (fst (List.split l))),
    hereditarily_pbarred T P l Hl Hl' -> Qind l Hl Hl'.
Proof.
  specialize (hereditarily_pbarred_ind _ _ T P Qind).
  intros Hind l Hl Hl' t.
  destruct t as [[HPl | (HPl & Ht)]].
  - apply (fst (Hind l Hl Hl')).
    assumption.
  - apply (snd (Hind l Hl Hl')); auto.
Qed.

Fixpoint hereditarily_pbarred_rect [Q A : Type] (T : sigma1_tree Q A) (P : list (Q * A) -> Prop)
  (Hdec : forall (l : list (Q * A)), T.(ptree_path) l -> {P l} + {~ P l})
  (Qind : forall (l : list (Q * A)), T.(ptree_path) l -> (List.NoDup (fst (List.split l))) -> Type) :
  (forall (l : list (Q * A)) (Hl : T.(ptree_path) l) (Hl' : List.NoDup (fst (List.split l))),
    (P l -> Qind l Hl Hl') * (
      ~ P l ->
      (forall (q : Q) (a : A) (Hal : T.(ptree_path) ((q, a) :: l))
        (Hq : ~ List.In q (fst (List.split l))),
          hereditarily_pbarred T P ((q, a) :: l) Hal (cons_no_dup _ _ _ Hq Hl')) ->
      (forall (q : Q) (a : A) (Hal : T.(ptree_path) ((q, a) :: l))
        (Hq : ~ List.In q (fst (List.split l))),
          Qind ((q, a) :: l) Hal (cons_no_dup _ _ _ Hq Hl')) ->
    Qind l Hl Hl')
  ) -> forall (l : list (Q * A)) (Hl : T.(ptree_path) l) (Hl' : List.NoDup (fst (List.split l))),
    hereditarily_pbarred T P l Hl Hl' -> Qind l Hl Hl'.
Proof.
  specialize (hereditarily_pbarred_rect _ _ _ _ Hdec Qind).
  intros Hind l Hl Hl' t.
  destruct t as [t].
  destruct (Hdec l Hl) as [HPl | HPl].
  - apply (fst (Hind l Hl Hl')).
    assumption.
  - assert (Hhered : forall (q : Q) (a : A) (Hal : ptree_path T ((q, a) :: l))
    (Hq : ~ List.In q (fst (List.split l))),
    hereditarily_pbarred T P ((q, a) :: l) Hal (cons_no_dup _ _ _ Hq Hl'))
    by (destruct t as [HPl' | (_ & Hhered)]; [contradiction | auto]).
    apply (snd (Hind l Hl Hl')); auto.
Defined.

Definition pBIT (Q A : Type) : Prop :=
  forall (T : sigma1_tree Q A) (P : list (Q * A) -> Prop),
    (forall (l : list (Q * A)), T.(ptree_path) l -> {P l} + {~P l}) ->
    (forall (l l' : list (Q * A)), P l -> P (l' ++ l)) ->
    pbarred T P -> forall (Hnil : T.(ptree_path) nil) (Hnil' : List.NoDup nil),
      hereditarily_pbarred T P nil Hnil Hnil'.

Lemma generalize_pBIT_iff_pBIT : forall (A : Type),
  pBIT nat A <-> @Neighborhood.pBIT delay_monad A.
Proof.
  intros A.
  split.
  - intros HPBIT T P Hdec Hmono Hbar Hnil.
    (*
    P' (l : list (nat * A)) := pour tout i, fst li = i et P (snd (split l))
    T' : construit à partir de T en transformant la fonction de voisinage en fonction d'arbre
    P' est décidable sur T et monotone et barré sur T, donc inductivement barré sur T
    je pense qu'il est aussi barré sur T
    *)

    admit. (* vrai *)
  - intros HBPIT T P Hdec Hmono Hbar Hnil.
    (*
    P' (l : list nat) := P (combine (iota (length l)) l)
    T' : c'est là tout le problème. T' est un arbre contrôlé par un "voisinage".
    Comment le transformer en arbre contrôlé par un "arbre extensionnel" ?
    *)

    admit. (* faux *)
Abort.

(* TODO: export *)
Lemma mem_fst_dec [X Y : Type] {Xdec : forall (x y : X), {x = y} + {x <> y}} :
  forall (l : list (X * Y)) (x : X),
    {List.In x (fst (List.split l))} + {~List.In x (fst (List.split l))}.
Proof.
  intros l x.
  apply List.In_dec.
  apply Xdec.
Defined.

Definition mem_snd_in [X Y : Type] {Xdec : forall (x y : X), {x = y} + {x <> y}} :
  forall (l : list (X * Y)) (x : X), List.In x (fst (List.split l)) ->
    { y : Y | List.In (x, y) l }.
Proof.
  intros l x' Hx'.
  induction l as [|(x, y) l IHl].
  - inversion Hx'.
  - simpl in *.
    destruct (List.split l).
    simpl in *.
    destruct (Xdec x x') as [<- | E].
    + exists y.
      left.
      reflexivity.
    + assert (HIn : List.In x' l0)
      by (destruct Hx'; [contradiction | assumption]).
      apply IHl in HIn as (y' & Hy').
      exists y'.
      right; assumption.
Defined.

Fixpoint fold_part_list_general [Q A R : Type] (tau : @ext_tree delay_monad Q A R) (f : Q -> delay A) (n : nat) : delay (list (Q * A)) :=
  match n with
  | O => Now []
  | S n => bind (fold_part_list_general tau f n) (fun l =>
    bind (tau (snd (List.split l))) (fun o => match o with
    | output r => Now l
    | ask q => bind (f q) (fun a => Now ((q, a) :: l))
    end))
  end.

Lemma fold_part_list_general_hasvalue [Q A R : Type] (tau : @ext_tree delay_monad Q A R) (f : Q -> delay A) :
  forall (n : nat) (l : list (Q * A)),
    fold_part_list_general tau f n =! l ->
      List.Forall (fun '(q, a) => f q =! a) l.
Proof.
  intros n.
  induction n as [|n IHn]; intros l Hl; cbn in *.
  - inversion Hl; constructor.
  - rename l into l'.
    apply bind_hasvalue in Hl as (l & Hl & ([q | r] & Ho & Hl')%bind_hasvalue).
    + apply bind_hasvalue in Hl' as (a & Ha & Hl').
      inversion Hl'; subst.
      constructor.
      1: assumption.
      apply IHn.
      assumption.
    + inversion Hl'; subst.
      apply IHn.
      assumption.
Qed.

(* TODO : export *)
Fixpoint part_fold `{Partiality : partiality} [A : Type]
  (f : nat -> A -> part A) (acc : A) (n : nat) : part A
:=
  match n with
  | O => partial.ret acc
  | S n => partial.bind (part_fold f acc n) (fun acc => f n acc)
  end.

Lemma part_fold_ter `{Partiality : partiality} [A : Type]
  (f : nat -> A -> part A) (acc : A) :
  forall (k n : nat), k <= n -> ter (part_fold f acc n) -> ter (part_fold f acc k).
Proof.
  intros k n Hk.
  revert acc.
  induction Hk as [| n Hk IHk ]; intros acc Hn.
  1: assumption.
  cbn in Hn.
  destruct Hn as (_ & (r & Hn & _)%partial.bind_hasvalue).
  apply IHk.
  eexists; eassumption.
Qed.

Lemma part_fold_preserves `{Partiality : partiality} [A : Type]
  (f : nat -> A -> part A) (acc : A) (k n : nat) (a : A) :
  k <= n -> part_fold f acc k =! a ->
  (forall (k : nat), k <= k -> k <= n -> f k a =! a) ->
  part_fold f acc n =! a.
Proof.
  intros Hkn.
  induction Hkn as [| n Hkn IHkn ]; intros Hk Ha.
  1: assumption.
  cbn.
  rewrite bind_hasvalue_given.
  1: apply Ha; lia.
  apply IHkn; auto.
Qed.

(* TODO : export, mais aussi : sert à rien *)
Lemma without_loss_of_generality_le_sym (P : nat -> nat -> Prop) :
  (forall (n m : nat), n <= m -> P n m) ->
  ((forall (n m : nat), P n m <-> P m n)) ->
  forall (n m : nat), P n m.
Proof.
  intros Hle Hsym n m.
  destruct (le_ge_dec n m) as [Hnm | Hmn].
  - apply Hle, Hnm.
  - apply Hsym, Hle, Hmn.
Qed.

Definition mu_fold `{Partiality : partiality} [A B : Type]
  (f : nat -> A -> part (result (Q := A) (R := B))) (acc : A) : part B
:=
  mu_opt (fun n => partial.bind (part_fold (fun n o => match o with
    | output r => partial.ret (output r)
    | ask a => f n a
    end) (ask acc) n) (fun r => match r with
  | output r => partial.ret (Some r)
  | ask _ => partial.ret (None)
  end)).

Lemma constructive_minimizer (P : nat -> Prop) :
  (exists (n : nat), P n /\
    inhabited (forall (k : nat), k <= n -> {P k} + {~ P k})) ->
  (exists (n : nat), P n /\ forall (k : nat), k < n -> ~ P k).
Proof.
  intros (n & Hn & (Hdec)).
  set (Q k := if le_gt_dec k n then P k else False).
  assert (HQdec : forall (k : nat), {Q k} + {~ Q k}). {
    intros k.
    unfold Q.
    destruct (le_gt_dec k n) as [Hk | Hk].
    - apply Hdec; assumption.
    - right; exact id.
  }
  assert (Hn' : exists (n : nat), Q n). {
    exists n.
    unfold Q.
    destruct le_gt_dec as [Hn' | Hn'].
    - assumption.
    - apply Arith_base.gt_irrefl_stt in Hn'.
      assumption.
  }
  apply ConstructiveEpsilon.epsilon_smallest in Hn' as (m & Hm & Hm').
  2: assumption.
  exists m.
  unfold Q in Hm, Hm'.
  destruct (le_gt_dec m n); try now inversion Hm.
  split.
  1: assumption.
  intros k Hk Hc.
  assert (Hk' : k <= n) by lia.
  contradict Hk.
  apply Nat.nlt_ge.
  apply Hm'.
  destruct (le_gt_dec k n) as [Hk | Hk].
  1: assumption.
  lia.
Qed.

Lemma mu_fold_spec `{Partiality : partiality} [A B : Type]
  (f : nat -> A -> part (result (Q := A) (R := B))) (acc : A) :
  forall (b : B),
    mu_fold f acc =! b <->
      exists (n : nat), part_fold (fun n o => match o with
      | output r => partial.ret (output r)
      | ask a => f n a
      end) (ask acc) n =! output b.
Proof.
  intros b.
  split.
  - intros (n & ([|] & Hr & [=]%ret_hasvalue_iff)%partial.bind_hasvalue & Hn')%mu_opt_hasvalue; subst.
    exists n.
    assumption.
  - eset (P n := part_fold _ (ask acc) n =! output b).
    intros Hb.
    assert (exists (n : nat), P n /\ forall (k : nat), k < n -> ~ P k)
    as (n & Hn & Hn'). {
      apply constructive_minimizer.
      destruct Hb as (n & Hn).
      exists n.
      split.
      1: assumption.
      constructor.
      intros k Hk.
      unfold P in *.
      destruct (eval (part_fold_ter _ _ _ _ Hk ltac:(eexists; eassumption))) eqn:E.
      - right.
        intros Hc.
        enough (output b = ask a) as [=].
        eapply partial.hasvalue_det.
        + apply Hc.
        + rewrite <- E; apply eval_hasvalue.
      - left.
        eenough (output b0 = output b) as <-. {
          rewrite <- E.
          apply eval_hasvalue.
        }
        eapply partial.hasvalue_det.
        2: apply Hn.
        apply part_fold_preserves with (k := k); auto.
        2: intros; psimpl.
        rewrite <- E.
        apply eval_hasvalue.
    }
    unfold mu_fold.
    apply mu_opt_hasvalue.
    exists n.
    split.
    + unfold P in Hn.
      rewrite bind_hasvalue_given.
      2: eassumption.
      psimpl.
    + intros k Hk.
      specialize (Hn' k Hk).
      unfold P in Hn, Hn'.
      eassert (ter (part_fold _ _ _)) as (o & Ho). {
        eapply part_fold_ter with (k := k).
        2: eexists; eauto.
        apply Nat.lt_le_incl; assumption.
      }
      rewrite partial.bind_hasvalue_given.
      2: apply Ho.
      destruct o.
      1: psimpl.
      eenough (E : output b0 = output b)
      by (inversion E; subst; contradiction).
      eapply partial.hasvalue_det.
      2: eassumption.
      eapply part_fold_preserves.
      3: intros; psimpl.
      2: eassumption.
      apply Nat.lt_le_incl; assumption.
Qed.

(* _ :=
on regarde tau  : si ça renvoie ask q, on prend lcombine0 = (q, l0)
on regarde tau l[0] : si ça renvoie ask q, on prend lcombine1 = (q, l1)
etc, etc.
et si ça renvoie output _, on renvoie ce qu'on avait jusque là
*)

Fixpoint rebuild_first_values `{Partiality : partiality} [Q A R : Type]
  (tau : @ext_tree _ Q A R) (l : list A) : part (list (Q * A))
:=
  match l with
  | [] => partial.ret []
  | a :: l => partial.bind (tau l) (fun o => match o with
    | output r => rebuild_first_values tau l
    | ask q => partial.bind (rebuild_first_values tau l) (fun lcombine =>
      partial.ret ((q, a) :: lcombine))
    end)
  end.

Lemma rebuild_first_values_app `{Partiality : partiality} [Q A R : Type] (tau : @ext_tree _ Q A R) :
  forall (la1 la2 : list A) (lqa : list (Q * A)),
    rebuild_first_values tau (la2 ++ la1) =! lqa ->
    exists (lqa1 lqa2 : list (Q * A)),
      rebuild_first_values tau la1 =! lqa1 /\ lqa = lqa2 ++ lqa1.
Proof.
  intros la1 la2.
  induction la2 as [| a la2 IHla2 ]; intros lqa Hlqa; cbn in *.
  - exists lqa, [].
    split; auto.
  - apply partial.bind_hasvalue in Hlqa as ([q | r] & Hr & Hlqa).
    + apply partial.bind_hasvalue in Hlqa as (lqa' & (lqa1 & lqa2 & HLqa1 & ->)%IHla2 & <-%partial.ret_hasvalue_iff).
      eexists _, _.
      split; eauto.
      change ((q, a) :: lqa2 ++ lqa1) with (((q, a) :: lqa2) ++ lqa1).
      reflexivity.
    + apply IHla2 in Hlqa as (lqa1 & lqa2 & Hlqa1 & ->).
      eexists _, _.
      split; eauto.
Qed.

Definition normalize_ext_tree `{Partiality : partiality} [Q A R : Type] {Qdec : forall (x y : Q), {x = y} + {x <> y}}
  (tau : @ext_tree _ Q A R) : ext_tree
:=
  fun l => partial.bind (rebuild_first_values tau l) (fun lcombine =>
    mu_fold (fun (n : nat) (lcombine : list (Q * A)) =>
      partial.bind (tau (snd (List.split lcombine))) (fun o => match o with
      | output r => partial.ret (output (output r))
      | ask q => match mem_fst_dec (Xdec := Qdec) lcombine q with
        | left Hq => partial.ret (ask ((q, ` (mem_snd_in (Xdec := Qdec) lcombine q Hq)) :: lcombine))
        | right _ => partial.ret (output (ask q))
        end
      end)
    ) lcombine
  ).

Definition normalized_ext_tree `{Partiality : partiality} [Q A R : Type] (tau : @ext_tree _ Q A R) : Prop :=
  forall (l l' : list A) (q : Q),
    tau l =! ask q -> tau (l' ++ l) =! ask q -> l' = [].

Definition normalized_strict_well_founded_tree_fun_for `{Partiality : partiality} [Q A R : Type] (tau : @ext_tree _ Q A R)
  (F : (Q -> part A) -> part R) : Prop := normalized_ext_tree tau /\ strict_well_founded_tree_fun_for tau F.

Definition normalized_strict_well_founded_tree_fun_cont `{Partiality : partiality} [Q A R : Type] (F : (Q -> part A) -> part R) : Prop :=
  exists (tau : ext_tree), normalized_strict_well_founded_tree_fun_for tau F.

Lemma normalize_is_normalized `{Partiality : partiality} [Q A R : Type] {Qdec : forall (x y : Q), {x = y} + {x <> y}}
  (tau : @ext_tree _ Q A R) :
  normalized_ext_tree (normalize_ext_tree (Qdec := Qdec) tau).
Proof.
  intros l l' q (lcombine & Hlcombine & Hl)%partial.bind_hasvalue (l'combine & Hl'combine & Hl')%partial.bind_hasvalue.
  assert (exists l'combine', l'combine = l'combine' ++ lcombine) as (l'combine' & ->). {
    apply rebuild_first_values_app in Hl'combine as (lqa1 & lqa2 & Hlqa1 & Hlqa2).
    enough (lqa1 = lcombine) as <- by (eexists _; eauto).
    eapply partial.hasvalue_det; eauto.
  }
  rename l'combine' into l'combine.
  (* TODO : changer les combine en combines *)
  (* btw j'ai besoin que ce soit wf, au moins *)
  (* et peut-être aussi strict *)
Admitted.

(* TODO : meilleur nom *)
Unset Elimination Schemes.
Inductive good_list [Q A R : Type] (tau : @ext_tree _ Q A R) (l : list (Q * A)) : Prop :=
| GoodList : (forall (q : Q) (a : A) (l' : list (Q * A)), l = (q, a) :: l' ->
    good_list tau l' /\ tau (snd (List.split l')) =! ask q) -> good_list tau l.
Set Elimination Schemes.

(* TODO : rect *)
Fixpoint good_list_ind [Q A R : Type] (tau : @ext_tree _ Q A R) (P : list (Q * A) -> Prop)
  (Hnil : P [])
  (Hcons : forall (q : Q) (a : A) (l : list (Q * A)),
    tau (snd (List.split l)) =! ask q -> good_list tau l -> P l -> P ((q, a) :: l)) :
  forall (l : list (Q * A)), good_list tau l -> P l.
Proof.
  specialize good_list_ind with (tau := tau) (1 := Hnil) (2 := Hcons).
  intros l Hl.
  destruct l as [|(q, a) l].
  - apply Hnil.
  - destruct Hl as [Hl].
    specialize Hl with (1 := eq_refl) as (Hl' & Hl).
    apply Hcons; auto.
Defined.

Lemma good_list_extract [Q A R : Type] (tau : @ext_tree _ Q A R) (l : list (Q * A)) :
  good_list tau l ->
  forall (q : Q), List.In q (fst (List.split l)) ->
    exists (n : nat), tau (List.skipn (S n) (snd (List.split l))) =! ask q
      /\ List.nth_error (fst (List.split l)) n = Some q.
Proof.
  intros Hl q' Hq'.
  induction Hl as [| q a l Hl Hl' IHl].
  1: inversion Hq'.
  cbn [List.split] in *.
  destruct (List.split l) as (lq & la) eqn:E.
  apply List.split_combine in E as <-.
  cbn [fst snd] in *.
  destruct Hq' as [-> | (n & Hla & Hlq)%IHl].
  - exists 0.
    cbn [List.skipn List.nth_error].
    rewrite List.skipn_0.
    auto.
  - exists (S n).
    cbn [List.skipn List.nth_error].
    split; assumption.
Qed.

Lemma nth_error_split_lr {A B : Type} (l : list (A * B)) (n : nat) :
  (exists (a : A), List.nth_error (fst (List.split l)) n = Some a)
  <-> (exists (b : B), List.nth_error (snd (List.split l)) n = Some b).
Proof.
  revert l.
  induction n as [|n IHn]; intros l;
  (destruct l as [| (a, b) l];
  [| destruct (List.split l) as (la & lb) eqn:E]);
  cbn in *.
  - split; intros (_ & [=]).
  - rewrite E; cbn.
    split; intros _; eexists; reflexivity.
  - split; intros (_ & [=]).
  - specialize (IHn l).
    rewrite E in *; cbn in *.
    assumption.
Qed.

Lemma normalized_good_list_nodup [Q A R : Type] (tau : @ext_tree _ Q A R) (l : list (Q * A)) :
  normalized_ext_tree tau -> good_list tau l -> List.NoDup (fst (List.split l)).
Proof.
  intros Htau Hl.
  induction Hl as [| q a l Hl Hl' IHl].
  1: constructor.
  simpl.
  destruct (List.split l) as (lq & la) eqn:E.
  cbn [fst snd] in *.
  constructor.
  2: auto.
  change lq with (fst (lq, la)).
  rewrite <- E.
  intros Hq.
  destruct good_list_extract with (1 := Hl') (2 := Hq)
  as (n & Hla & Hlq).
  rewrite E in Hq, Hla, Hlq.
  cbn [fst snd] in *.
  assert (exists (a' : A), List.nth_error la n = Some a')
  as (a' & Ha'%List.firstn_skipn_middle). {
    change la with (snd (lq, la)).
    rewrite <- E.
    apply nth_error_split_lr.
    eexists.
    rewrite E.
    eassumption.
  }
  rewrite <- Ha' in Hl.
  assert (List.firstn n la ++ [a'] = nil). {
    eapply Htau.
    - apply Hla.
    - rewrite <- List.app_assoc.
      apply Hl.
  }
  destruct (List.firstn n la); inversion H.
Qed.

Definition dialogue_tree_for_at [Q A R : Type] {Qdec : forall (x y : Q), {x = y} + {x <> y}}
  (d : dialogue_tree Q A R) (F : (Q -> delay A) -> delay R) (l : list (Q * A)) : Prop
:= dialogue_tree_for d (fun f => F (fun q =>
  match mem_fst_dec (Xdec := Qdec) l q with
  | left Hq => ret (` (mem_snd_in (Xdec := Qdec) _ _ Hq))
  | right _ => f q
  end)).

Lemma good_list_extend_before [Q A R : Type] {Qdec : forall (x y : Q), {x = y} + {x <> y}}
  (tau : @ext_tree _ Q A R) (f : Q -> delay A) :
  forall (l : list (Q * A)), good_list tau l ->
    forall (n : nat), n < List.length l ->
    equiv
      (eval_ext_tree tau (fun q =>
        match mem_fst_dec (Xdec := Qdec) l q with
        | left Hq => ret (` (mem_snd_in (Xdec := Qdec) _ _ Hq))
        | right _ => f q
        end) n)
      (tau (List.skipn (List.length l - n) (snd (List.split l)))).
Proof.
  intros l Hl n.
  revert l Hl.
  induction n as [|n IHn]; intros l Hl Hn.
  - cbn.
    rewrite Nat.sub_0_r, <- List.length_snd_split, List.skipn_all.
    reflexivity.
  - unfold eval_ext_tree.
    rewrite eval_ext_tree_aux_succ.
    change (eval_ext_tree_aux ?tau ?f ?n [])
    with (eval_ext_tree tau f n).
    rewrite IHn.
    3: lia.
    2: assumption.
    assert (Hn' : length l - n = S (length l - S n)) by lia.
    rewrite Hn'.
    remember (length l - S n) as k eqn:E.
    assert (exists (a : A), List.nth_error (snd (List.split l)) k = Some a)
    as (a & Ha) by admit.
    assert (List.skipn k (snd (List.split l))
      = a :: List.skipn (S k) (snd (List.split l))) as -> by admit.
    intros res; split.
    + admit.
    + intros Hres.
      assert (ter (tau (List.skipn (S k) (snd (List.split l))))) as (res' & Hres'). {
        admit. (* tau wf *)
      }
      rewrite bind_hasvalue_given; eauto.
      destruct res' as [q | r].
      2: admit.
      (* tau vient de me sortir q et l est good,
      donc on est dans le premier cas : j'ai récupérer le a *)

      (* tree_fun_trace bidule me donne fst (List.split (List.firstn n l))
      le fold fonctionne sans souci, vu que ça appartient à la liste
      je finis par évaluer tau (a :: List.skipn (S k) (snd (List.split l)))
      c'est parfait *)

      (* mais ça va être une purge à montrer ;
      plutôt que passer par les lemmes existants, autant
      les redémontrer avec les bonnes idées (le prédicat good list) *)

      admit.
Admitted.

Lemma good_list_extend_after [Q A R : Type] {Qdec : forall (x y : Q), {x = y} + {x <> y}}
  (tau : @ext_tree _ Q A R) (f : Q -> delay A) :
  forall (l : list (Q * A)), good_list tau l ->
    forall (n : nat), equiv
      (eval_ext_tree tau (fun q =>
        match mem_fst_dec (Xdec := Qdec) l q with
        | left Hq => ret (` (mem_snd_in (Xdec := Qdec) _ _ Hq))
        | right _ => f q
        end) (n + List.length l))
    (eval_ext_tree_aux tau f n (snd (List.split l))).
Proof.
  intros l Hl n.
  revert l Hl.
  induction n as [| n IHn]; intros l Hl.
  1: cbn.
  1: admit.
  change (S n + length l) with (S (n + length l)).
  unfold eval_ext_tree in *.
  rewrite 2! eval_ext_tree_aux_succ.
  apply equiv_bind.
  1: apply IHn; assumption.
  intros [q | r] Hr Hr'.
  2: reflexivity.
Admitted.

Lemma hereditary [Q A R : Type] {Qdec : forall (x y : Q), {x = y} + {x <> y}}
  (tau : @ext_tree _ Q A R) (F : (Q -> part A) -> part R)
  (Htau : strict_well_founded_tree_fun_for tau F)
  (Htau' : normalized_ext_tree tau) :
  forall (l : list (Q * A)),
    (termination_tree F tau Htau).(ptree_path) l ->
    good_list tau l ->
    (forall (q : Q) (a : A),
      (termination_tree F tau Htau).(ptree_path) ((q, a) :: l) ->
      good_list tau ((q, a) :: l) ->
      { d | dialogue_tree_for_at (Qdec := Qdec) d F ((q, a) :: l) }) ->
    { d | dialogue_tree_for_at (Qdec := Qdec) d F l }.
Proof.
  intros l Hl Hl' IHl.
  cbn in Hl.
  destruct (eval Hl) as [q | r] eqn:E.
  - unshelve eexists.
    + apply (Ask q).
      intros a.
      apply (bind (proof_ter (tau (a :: snd (List.split l))))).
      intros Hal.
      apply Now.
      assert (Hal' : (termination_tree F tau Htau).(ptree_path) ((q, a) :: l)). {
        cbn in *.
        destruct (List.split l).
        cbn.
        assumption.
      }
      assert (Hq : good_list tau ((q, a) :: l)). {
        constructor.
        intros q' a' l' E'; symmetry in E'; inversion E'; subst.
        split; try assumption.
        rewrite <- E; apply eval_hasvalue.
      }
      apply (IHl q a Hal' Hq).
    + intros f.
      rewrite eval_dialogue_tree_ask.
      repeat change (bind ?a ?f) with (partial.bind a f).
      setoid_rewrite equiv_bind_mult.
      setoid_rewrite equiv_bind_ret.
      assert (Hmap : forall (A B : Type) (p q : delay A) (f : A -> B),
        (forall (a a' : A), f a = f a' -> a = a') ->
        equiv (partial.bind p (fun r => ret (f r))) (partial.bind q (fun r => ret (f r))) -> equiv p q
      ). {
        intros X Y x x' g Hinj Heq.
        intros r; split;
          intros Hr;
          eassert (partial.bind _ (fun r => ret (g r)) =! g r) as Hr'
          by (rewrite bind_hasvalue_given; [ constructor | eassumption ]);
          apply Heq in Hr' as (r' & Hr' & Hr'')%bind_hasvalue;
          inversion Hr''; apply Hinj in H; subst; assumption.
      }
      eapply Hmap with (f := output).
      1: intros a a' [=]; assumption.
      clear Hmap.
      rewrite equiv_bind_mult.
      setoid_rewrite equiv_bind_mult.

      (* TODO : c'est faux, c'est le cas seulement si F termine. sinon, faut utiliser mu *)
      eassert (HF : exists (n : nat), equiv (eval_ext_tree tau f n) (bind (F f) (fun r => ret (output r)))). {
        destruct Htau as (? & (? & Htau)).
        pose proof (tree_fun_for_good_def F tau) as (H & _).
        apply H; eauto.
        admit.
      }
      admit.
  - exists (Output r).
    unfold dialogue_tree_for_at, dialogue_tree_for.
    intros f.
    cbn.
    eapply equiv_hasvalue.
    1: constructor.
    apply Htau.
    clear IHl.
    assert (Hr : tau (snd (List.split l )) =! output r)
    by (rewrite <- E; apply eval_hasvalue).
    clear Hl E.
    exists (List.length l).
    apply good_list_extend_after with (n := 0).
    1: assumption.
    cbn.
    assumption.
Admitted.

Lemma pBIT_ext_tree_to_dialogue_tree [Q A R : Type] `{Qdec : forall (x y : Q), {x = y} + {x <> y}} :
  pBIT Q A -> (forall (tau : ext_tree) (F : (Q -> delay A) -> delay R),
    strict_well_founded_tree_fun_for tau F -> ter (tau []) -> { d : dialogue_tree Q A R | dialogue_tree_for_at (Qdec := Qdec) d F [] }).
Proof.
Abort.