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.
Import EmbedNatNotations PartialTactics ListNotations.

Set Bullet Behavior "Strict Subproofs".

Section PartialZoo.

Lemma Forall_app {A : Type} {P : A -> Prop} :
  forall (l1 l2 : list A), Forall P l1 -> Forall P l2 ->
  Forall P (l1 ++ l2).
Proof.
  intros l1 l2 Hl1 Hl2.
  induction Hl1 as [| a l1 Hl1 IHl1]; try assumption.
  constructor; assumption.
Qed.

Lemma Forall_app_l {A : Type} {P : A -> Prop} :
  forall (l1 l2 : list A), Forall P (l1 ++ l2) -> Forall P l1.
Proof.
  intros l1 l2 H.
  induction l1 as [| a l1 IHl1 ]; auto.
  inversion H; subst.
  constructor; firstorder.
Qed.

Lemma Forall_app_r {A : Type} {P : A -> Prop} :
  forall (l1 l2 : list A), Forall P (l1 ++ l2) -> Forall P l2.
Proof.
  intros l1 l2 H.
  induction l1 as [| a l1 IHl1]; auto.
  inversion H; subst.
  firstorder.
Qed.

Lemma Forall2_app {A B : Type} {P : A -> B -> Prop} :
  forall (la1 la2 : list A) (lb1 lb2 : list B), Forall2 P la1 lb1 -> Forall2 P la2 lb2 ->
  Forall2 P (la1 ++ la2) (lb1 ++ lb2).
Proof.
  intros la1 la2 lb1 lb2 Hl1 Hl2.
  induction Hl1 as [| a b la1 lb1 Hl1 IHl1]; try assumption.
  constructor; assumption.
Qed.

Lemma Forall2_app_l {A B : Type} {P : A -> B -> Prop} :
  forall (la1 la2 : list A) (lb1 lb2 : list B),
    length la1 = length lb1 ->
      Forall2 P (la1 ++ la2) (lb1 ++ lb2) -> Forall2 P la1 lb1.
Proof.
  intros la1.
  induction la1 as [| a la1 IHl1 ]; intros la2 lb1 lb2 Hlen H.
  - symmetry in Hlen.
    rewrite length_zero_iff_nil in Hlen.
    subst.
    constructor.
  - destruct lb1 as [|b lb1].
    1: inversion Hlen.
    simpl in H.
    inversion H; subst.
    constructor; eauto.
Qed.

Lemma Forall2_app_r {A B : Type} {P : A -> B -> Prop} :
  forall (la1 la2 : list A) (lb1 lb2 : list B),
    length la1 = length lb1 ->
      Forall2 P (la1 ++ la2) (lb1 ++ lb2) -> Forall2 P la2 lb2.
Proof.
  intros la1.
  induction la1 as [| a la1 IHl1 ]; intros la2 lb1 lb2 Hlen H.
  - symmetry in Hlen.
    rewrite length_zero_iff_nil in Hlen.
    subst.
    simpl in H; eauto.
  - destruct lb1 as [|b lb1].
    1: inversion Hlen.
    simpl in H.
    inversion H; subst.
    eauto.
Qed.

Lemma Forall2_imp {A B : Type} {P Q : A -> B -> Prop} :
  forall (la : list A) (lb : list B), Forall2 P la lb ->
    Forall2 (fun a b => P a b -> Q a b) la lb -> Forall2 Q la lb.
Proof.
  intros la lb HP.
  induction HP as [| a b la lb Hab Hl IHl]; intros Himp.
  1: constructor.
  inversion Himp; subst.
  constructor; eauto.
Qed.

Lemma Forall2_rev {A B : Type} {P : A -> B -> Prop} :
  forall (la : list A) (lb : list B),
    Forall2 P la lb -> Forall2 P (rev la) (rev lb).
Proof.
  intros la lb Hl.
  induction Hl as [|a b la lb Hab Hl IHl]; auto.
  simpl.
  apply Forall2_app; auto.
Qed.

Add Parametric Morphism (A B : Type) : ( @Forall2 A B )
  with signature pointwise_relation _ (pointwise_relation _ iff)
    ==> eq
    ==> eq
    ==> iff
  as Forall2_morph.
Proof.
  intros P Q H la lb.
  split; induction 1; constructor; try apply H; eauto.
Qed.

Lemma nth_error_app [A : Type] : forall (n : nat) (l l' : list A) (a : A),
  nth_error l n = Some a -> nth_error (l ++ l') n = Some a.
Proof.
  intros n.
  induction n as [|n IHn]; intros l l' a Ha;
  simpl in *; destruct l; simpl; inversion Ha; auto.
  rewrite Ha.
  eauto.
Qed.

Lemma firstn_app_l [A : Type] : forall (n : nat) (l l' : list A),
  n <= length l -> firstn n l = firstn n (l ++ l').
Proof.
  intros n.
  induction n as [|n IHn]; intros l l' Hn.
  1: reflexivity.
  destruct l as [|a l].
  1: inversion Hn.
  simpl in Hn.
  apply le_S_n in Hn.
  cbn.
  f_equal.
  auto.
Qed.

Lemma rev_firstn_rev [A : Type] : forall (n : nat) (l : list A),
  rev (skipn n (rev l)) ++ rev (firstn n (rev l)) = l.
Proof.
  intros n l.
  rewrite <- rev_involutive.
  generalize dependent (rev l); clear l; intros l.
  rewrite <- rev_app_distr.
  f_equal.
  revert l.
  induction n as [|n IHn]; intros l.
  1: reflexivity.
  destruct l as [|a l].
  1: reflexivity.
  simpl.
  f_equal.
  apply IHn.
Qed.

Context `{Partiality : partiality}.

#[global]
Add Parametric Relation A : (part A) ( @equiv _ A ) as equiv.

#[global]
Add Parametric Morphism A B : ( @bind _ A B )
  with signature @equiv _ A
    ==> pointwise_relation _ ( @equiv _ B )
    ==> @equiv _ B
  as bind_morph.
Proof.
  intros p q H f g H'.
  intros v.
  split; intros Hv; psimpl;
  try (apply H || apply H');
  eassumption.
Qed.

#[global]
Add Parametric Morphism A : ( @hasvalue _ A )
  with signature @equiv _ A
    ==> @eq A
    ==> iff
  as hasvalue_morph.
Proof.
  intros p q H.
  apply H.
Qed.

#[global]
Add Parametric Morphism A : ( @ter _ A )
  with signature @equiv _ A
    ==> iff
  as ter_morph.
Proof.
  intros p q H.
  split; intros (a & Ha);
  eexists; eapply H; eauto.
Qed.

Lemma bind_ter {A B : Type} :
  forall (p : part A) (f : A -> part B) (a : A),
  hasvalue p a -> ter (f a) -> ter (bind p f).
Proof.
  intros p f a Ha (b & Hb).
  exists b.
  psimpl.
Qed.

Lemma bind_ter' {A B : Type} :
  forall (p : part A) (f : A -> part B),
  ter p -> (forall (a : A), hasvalue p a -> ter (f a)) -> ter (bind p f).
Proof.
  intros p f (a & Ha) Hp.
  destruct (Hp a Ha).
  exists x.
  psimpl.
Qed.

Definition proof_ter {A : Type} (p : part A) : part (ter p).
Proof.
  set (p' := mu_tot (fun n => match seval p n with
  | Some _ => true
  | None => false
  end)).
  apply (bind p').
  intros n.
  destruct (seval p n) eqn:E.
  - apply ret.
    exists a.
    rewrite seval_hasvalue.
    eauto.
  - apply undef.
Defined.

Definition proof_ter_spec {A : Type} :
  forall (p : part A), ter p -> ter (proof_ter p).
Proof.
  intros p Hp.
  destruct Hp as (a & Ha).
  unfold proof_ter.
  assert (Hp : ter (mu_tot (fun n : nat => match seval (A:=A) p n with
    | Some _ => true
    | None => false
    end))). {
    rewrite seval_hasvalue in Ha.
    destruct Ha as (n & Hn).
    apply mu_tot_ter with n.
    rewrite Hn.
    reflexivity.
  }
  destruct Hp as (n & Hn).
  eapply bind_ter;
  psimpl.
  assert (Ha' : seval p n = Some a). {
    rewrite mu_tot_hasvalue in Hn.
    destruct Hn as (Hn & _).
    destruct (seval p n) as [a'|] eqn:E.
    2: inversion Hn.
    f_equal.
    assert (Ha0 : hasvalue p a'). {
      rewrite seval_hasvalue.
      eauto.
    }
    eapply hasvalue_det; eassumption.
  }
  pose (body (a' : A) (E : seval p n = Some a') :=
    ex_intro (fun a1 : A => p =! a1) a'
    (Morphisms.flip2 Morphisms.iff_flip_impl_subrelation (p =! a')
    (exists n0 : nat, seval (A:=A) p n0 = Some a')
    (seval_hasvalue p a')
    (ex_intro (fun n0 : nat => seval (A:=A) p n0 = Some a') n E))).
  change (ter (
    match seval p n as o return ((seval p n = o) ↛ (ter p))
    with
    | Some a' => fun E => ret (body a' E)
    | None => fun _ => undef
    end eq_refl)
  ).
  generalize dependent body.
  intros body.
  exists (body a Ha').
  destruct (seval p n).
  2: inversion Ha'.
  dependent elimination Ha'.
  psimpl.
Qed.

Lemma equiv_bind [A B : Type] (p p' : part A) (f f' : A -> part B) :
  equiv p p' ->
  (forall (a : A), hasvalue p a -> hasvalue p' a -> equiv (f a) (f' a)) ->
  equiv (bind p f) (bind p' f').
Proof.
  intros H Ha.
  assert (Ha' : forall a, p =! a -> equiv (f a) (f' a))
  by (intros; eapply Ha; try rewrite <- H; eauto).
  rewrite <- H.
  intros v; split; intros Hv; psimpl; (rewrite Ha' || rewrite <- Ha'); eauto.
Qed.

Lemma equiv_bind' [A B : Type] (p p' : part A) (f f' : A -> part B) :
  forall (a : A), hasvalue p a -> hasvalue p' a -> equiv (f a) (f' a) ->
  equiv (bind p f) (bind p' f').
Proof.
  intros a Hp Hp' H b.
  split; intros Hb; psimpl;
  assert (a = x) by (eapply hasvalue_det;
    [ | eassumption ]; eassumption);
  subst; apply H; assumption.
Qed.

Lemma equiv_hasvalue [A : Type] (p p' : part A) :
  forall (a : A), hasvalue p a -> hasvalue p' a ->
    equiv p p'.
Proof.
  intros a Hp Hp' a'; split; intros H;
  assert (a = a')
    by (eapply hasvalue_det;
    [ | eassumption ]; eassumption);
    subst; eassumption.
Qed.

Lemma equiv_bind_mult [A B C : Type] (p : part A) (f : A -> part B) (g : B -> part C) :
  equiv (bind (bind p f) g) (bind p (fun a => bind (f a) g)).
Proof.
  intros v; split; intros; psimpl.
Qed.

Lemma equiv_bind_ret [A B : Type] (a : A) (f : A -> part B) :
  equiv (bind (ret a) f) (f a).
Proof.
  intros v; split; intros; psimpl.
Qed.

Lemma equiv_ret_bind [A B : Type] (p : part A) :
  equiv (bind p (ret (A := A))) p.
Proof.
  intros v; split; intros; psimpl.
Qed.

Lemma equiv_bind_hasvalue [A B : Type] (p : part A) (f : A -> part B) (a : A) :
  p =! a -> equiv (bind p f) (f a).
Proof.
  intros H.
  rewrite <- (equiv_bind_ret a f).
  apply bind_morph; try reflexivity.
  eapply equiv_hasvalue; psimpl.
Qed.

Fixpoint tabulate_aux [A : Type] (f : nat -> A) (m n k : nat) : list A :=
  match k with
  | O => nil
  | S k => f m :: tabulate_aux f (S m) n k
  end.

Definition tabulate [A : Type] (f : nat -> A) (m n : nat) : list A := tabulate_aux f m n (n - m).

Lemma tabulate_same [A : Type] : forall (f : nat -> A) (n : nat),
  tabulate f n n = nil.
Proof.
  intros f n.
  unfold tabulate.
  rewrite Nat.sub_diag.
  reflexivity.
Qed.

Lemma tabulate_succ [A : Type] : forall (f : nat -> A) (m n : nat),
  S m <= n -> tabulate f m n = f m :: tabulate f (S m) n.
Proof.
  intros f m n H.
  unfold tabulate.
  rewrite Nat.sub_succ_r.
  set (k := n - m).
  assert (Hk : k <> 0) by lia.
  destruct k as [|k].
  1: contradiction Hk; reflexivity.
  reflexivity.
Qed.

Lemma tabulate_succ' [A : Type] : forall (f : nat -> A) (m n : nat),
  m <= n -> tabulate f m (S n) = tabulate f m n ++ [f n].
Proof.
  intros f m n H.
  unfold tabulate.
  rewrite Nat.sub_succ_l; try lia.
  remember (n - m) as k.
  generalize dependent n.
  generalize dependent m.
  induction k; intros m n H Hk.
  - simpl.
    assert (n = m) by lia.
    subst.
    reflexivity.
  - cbn.
    assert (Hk' : k = n - S m) by lia.
    rewrite Hk' at 1.
    assert (E : tabulate_aux f (S (S m)) (S n) (n - S m) = tabulate f (S (S m)) (S n))
    by reflexivity.
    rewrite E.
    rewrite <- tabulate_succ; try lia.
    f_equal.
    rewrite <- IHk; try lia.
    cbn.
    rewrite <- Hk.
    cbn.
    reflexivity.
Qed.

Lemma tabulate_add [A : Type] : forall (f : nat -> A) (m n k : nat),
  m <= n -> n <= k ->
  tabulate f m n ++ tabulate f n k = tabulate f m k.
Proof.
  intros f m n k Hm.
  revert k.
  induction Hm as [| n Hm IHm]; intros k Hk.
  - rewrite tabulate_same.
    simpl.
    reflexivity.
  - rewrite tabulate_succ'; try eassumption.
    rewrite <- app_assoc.
    simpl.
    rewrite <- tabulate_succ; try eassumption.
    apply IHm; lia.
Qed.

Lemma tabulate_length [A : Type] : forall (f : nat -> A) (m n k : nat),
  length (tabulate_aux f m n k) = k.
Proof.
  intros f m n k. revert m n.
  induction k as [|k IHk]; intros m n.
  - reflexivity.
  - simpl.
    rewrite IHk.
    reflexivity.
Qed.

Lemma tabulate_map [A B : Type] : forall (f : nat -> A) (g : A -> B) (m n k : nat),
  map g (tabulate_aux f m n k) = tabulate_aux (fun n => g (f n)) m n k.
Proof.
  intros f g m n k.
  revert m n.
  induction k as [|k IHk]; intros m n; cbn.
  - reflexivity.
  - f_equal.
    apply IHk.
Qed.

Lemma tabulate_equiv [A : Type] : forall (f g : nat -> part A) (n : nat),
  (forall (k : nat), k < n -> equiv (f k) (g k)) ->
    Forall2 (fun p q => equiv p q) (tabulate f 0 n) (tabulate g 0 n).
Proof.
  intros f g n Hn.
  rewrite <- rev_involutive with (l := tabulate f 0 n), <- rev_involutive with (l := tabulate g 0 n), Forall2_rev.
  1: {
    apply list_relations.Forall_Forall2_diag.
    apply list_relations.Forall_true.
    reflexivity.
  }
  induction n as [|n IHn].
  1: cbn; constructor.
  specialize (IHn (fun k Hk => Hn k ltac:(eauto with arith))).
  do 2 (rewrite tabulate_succ', rev_unit; eauto with arith).
Qed.

Fixpoint fold_part_list_aux [A : Type] (l : list (part A)) (l' : list A) : part (list A) :=
  match l with
  | nil => ret l'
  | p :: l => bind p (fun a => fold_part_list_aux l (a :: l'))
  end.

#[global]
Add Parametric Morphism A : ( @fold_part_list_aux A )
  with signature Forall2 (equiv (A := A))
    ==> eq
    ==> equiv (A := list A)
  as fold_part_list_aux_morph.
Proof.
  intros l1 l2 Hl.
  induction Hl as [|p1 p2 l1 l2 Hp Hl IH]; intros l'; cbn.
  - reflexivity.
  - rewrite Hp.
    setoid_rewrite IH.
    reflexivity.
Qed.

Definition fold_part_list [A : Type] (l : list (part A)) : part (list A)
  := fold_part_list_aux l nil.

#[global]
Add Parametric Morphism A : ( @fold_part_list A )
  with signature Forall2 (equiv (A := A))
    ==> equiv (A := list A)
  as fold_part_list_morph.
Proof.
  intros l l' Hl.
  unfold fold_part_list.
  eapply fold_part_list_aux_morph; auto.
Qed.

Lemma fold_part_list_aux_app [A : Type] (l : list (part A)) (l1 l2 : list A) :
  forall (lv : list A), fold_part_list_aux l l1 =! lv ->
    fold_part_list_aux l (l1 ++ l2) =! (lv ++ l2).
Proof.
  revert l1 l2.
  induction l; cbn in *; intros l1 l2 lv Hlv; psimpl.
  change (x :: l1 ++ l2) with ((x :: l1) ++ l2).
  apply IHl.
  assumption.
Qed.

Lemma fold_part_list_aux_app' [A : Type] (l : list (part A)) (l1 l2 : list A) :
  forall (lv : list A), fold_part_list_aux l (l1 ++ l2) =! lv ->
    exists (lv' : list A), fold_part_list_aux l l1 =! lv' /\ lv = lv' ++ l2.
Proof.
  revert l1 l2.
  induction l; intros l1 l2 lv Hlv.
  - exists l1.
    cbn in *.
    split; psimpl.
  - cbn in Hlv.
    apply bind_hasvalue in Hlv as (x & Hx & Hlv).
    specialize (IHl (x :: l1) l2 lv Hlv).
    destruct IHl as (lv' & Hlv1 & Hlv2).
    exists lv'.
    split; eauto.
    cbn.
    psimpl.
Qed.

Lemma fold_part_list_cons [A : Type] (l : list (part A)) (p : part A) :
  equiv
    (fold_part_list (p :: l))
    (bind p (fun a =>
      bind (fold_part_list l) (fun l => ret (l ++ [a])))).
Proof.
  enough (H : forall l', equiv
    (fold_part_list_aux (p :: l) l')
    (bind p (fun a =>
      bind (fold_part_list l) (fun l => ret (l ++ (a :: l')))))
  ) by apply H.
  intros l'.
  induction l as [| q l IHl]; cbn.
  - apply bind_morph; try reflexivity.
    intros a.
    eapply equiv_hasvalue; psimpl.
  - apply equiv_bind; try reflexivity.
    intros a Ha _.
    rewrite equiv_bind_mult.
    apply equiv_bind; try reflexivity.
    intros a' Ha' _.
    cbn in IHl.
    rewrite equiv_bind_hasvalue in IHl; eauto.
    rewrite equiv_bind_hasvalue in IHl; eauto.
    intros v; split; intros Hv.
    + change (a' :: a :: l') with ([a'] ++ (a :: l')) in Hv.
      apply fold_part_list_aux_app' in Hv.
      destruct Hv as (lv & Hlv1 & Hlv2).
      subst.
      psimpl.
    + psimpl.
      eapply fold_part_list_aux_app in H.
      eapply H.
Qed.

Lemma fold_part_list_length [A : Type] (lp : list (part A)) (l' la : list A) :
  fold_part_list_aux lp l' =! la -> length lp + length l' = length la.
Proof.
  revert l' la.
  induction lp as [| p lp IHlp]; intros l' la H.
  - cbn in H.
    apply ret_hasvalue_inv in H.
    subst.
    reflexivity.
  - cbn in H.
    psimpl.
    apply IHlp in H0.
    cbn in *.
    rewrite <- H0; lia.
Qed.

Lemma fold_part_list_tabulate_add_length [A : Type] (f : nat -> part A) :
  forall (n m : nat) (l : list A),
    fold_part_list (tabulate f m (n + m)) =! l -> n = length l.
Proof.
  intros n m l Hn.
  apply fold_part_list_length in Hn.
  unfold tabulate in Hn.
  rewrite tabulate_length in Hn.
  now rewrite Nat.add_0_r, Nat.add_sub in Hn.
Qed.

Lemma fold_part_list_tabulate_length [A : Type] (f : nat -> part A) :
  forall (n : nat) (l : list A),
    fold_part_list (tabulate f 0 n) =! l -> n = length l.
Proof.
  intros n l Hn.
  eapply fold_part_list_tabulate_add_length.
  rewrite Nat.add_0_r.
  eauto.
Qed.

Lemma fold_part_list_app [A : Type] (l l' : list (part A)) :
  equiv (fold_part_list (l ++ l'))
    (bind (fold_part_list l) (fun l => bind (fold_part_list l')
      (fun l' => ret (l' ++ l)))).
Proof.
  induction l as [| p l IHl].
  - cbn.
    rewrite equiv_bind_ret.
    simpl.
    intros v; split; intros; psimpl; rewrite app_nil_r; psimpl.
  - change ((p :: l) ++ l') with (p :: (l ++ l')).
    do 2 rewrite fold_part_list_cons.
    intros v; split; intros Hv.
    + simpl_assms.
      rewrite IHl in H0.
      psimpl.
      rewrite <- app_assoc.
      psimpl.
    + simpl_assms.
      setoid_rewrite IHl.
      psimpl.
      rewrite <- app_assoc.
      psimpl.
Qed.

Lemma fold_part_list_ter [A : Type] : forall (l : list (part A)),
  ter (fold_part_list l) <-> Forall (fun p => ter p) l.
Proof.
  intros l.
  induction l as [|p l IHl]; split.
  1-2: intros _; try eexists; psimpl.
  - intros (la & Hla).
    rewrite fold_part_list_cons in Hla.
    psimpl.
    constructor.
    + eexists; eauto.
    + apply IHl; eexists; eauto.
  - intros Hl.
    inversion Hl; subst.
    rewrite fold_part_list_cons.
    eapply bind_ter'; eauto.
    intros a Ha.
    eapply bind_ter'.
    2: intros; eexists; psimpl.
    apply IHl; auto.
Qed.

Lemma fold_part_list_ret [A : Type] : forall (l : list A),
  fold_part_list (map (fun a => ret a) (rev l)) =! l.
Proof.
  induction l as [|a l IHl].
  1:psimpl.
  simpl.
  rewrite map_last, fold_part_list_app.
  psimpl.
Qed.

Definition mu_opt [A : Type] (f : nat -> part (option A)) : part A :=
  bind (mu (fun n => bind (f n) (fun o => ret (if o then true else false))))
    (fun n => bind (f n) (fun o => match o with | Some a => ret a | None => undef end)).

Lemma mu_opt_ter [A : Type] : forall (f : nat -> part (option A)),
  (exists (n : nat), (exists (a : A), f n =! Some a) /\ forall (k : nat), k < n -> f k =! None)
  -> ter (mu_opt f).
Proof.
  intros f (n & (a & Ha) & Hn).
  unfold mu_opt.
  eexists; rewrite bind_hasvalue; eexists; split; eauto.
  apply mu_hasvalue.
  split.
  1: rewrite bind_hasvalue; eexists; split; psimpl.
  intros m Hm.
  specialize (Hn m Hm).
  rewrite bind_hasvalue; eexists; split; psimpl.
  psimpl.
Qed.

Lemma mu_opt_hasvalue [A : Type] :
  forall (f : nat -> part (option A)) (a : A),
    (mu_opt f =! a) <-> (exists (n : nat),
      f n =! Some a /\ forall (k : nat), k < n -> f k =! None).
Proof.
  intros f a.
  split.
  - intros Ha.
    unfold mu_opt in *.
    psimpl.
    apply mu_hasvalue in H as (? & ?).
    destruct x0; psimpl.
    eexists; split; eauto.
    intros k Hk.
    specialize (H2 k Hk).
    psimpl.
    destruct x1; inversion H6.
    auto.
  - intros (n & Hn & Hf).
    unfold mu_opt.
    psimpl.
    apply mu_hasvalue.
    split.
    1: rewrite bind_hasvalue; eexists; split; psimpl.
    intros k Hk.
    specialize (Hf k Hk).
    psimpl.
Qed.

(* l = f (n - 1); ...; f 0 *)
Definition extend [A : Type] (l : list A) (f : nat -> part A) :=
  Forall2 (fun n a => f n =! a) (tabulate id 0 (length l)) (rev l).

Lemma extend_fold_part {A : Type} : forall (f : nat -> part A) (l : list A),
  extend l f <-> fold_part_list (tabulate f 0 (length l)) =! l.
Proof.
  intros f l.
  induction l as [|a l IHl].
  - split; intros; psimpl.
    constructor.
  - split; unfold extend; simpl; intros.
    all: rewrite tabulate_succ' in *; try lia.
    all: rewrite fold_part_list_app in *.
    + pose proof (H' := H).
      assert (length (tabulate id 0 (length l)) = length (rev l)). {
        rewrite length_rev.
        unfold tabulate.
        rewrite tabulate_length.
        apply Nat.sub_0_r.
      }
      apply Forall2_app_l in H; auto.
      apply Forall2_app_r in H'; auto.
      inversion H'; subst.
      clear H6 H0 H'.
      apply IHl in H.
      psimpl.
    + cbn in H.
      psimpl.
      apply IHl in H.
      apply Forall2_app; eauto.
Qed.

Lemma extend_app {A : Type} : forall (f : nat -> part A) (l l' : list A),
  extend (l' ++ l) f -> extend l f.
Proof.
  intros f l l' Hl.
  rewrite extend_fold_part in Hl |- *.
  rewrite length_app, <- tabulate_add with (n := length l),
    fold_part_list_app in Hl; try lia.
  psimpl.
  enough (x = l) by (subst; eauto).
  clear H.
  apply fold_part_list_tabulate_add_length in H0.
  generalize dependent x0.
  induction l' as [| a l' IHl']; intros x0 Hx0 Hl.
  1: symmetry in Hx0; apply length_zero_iff_nil in Hx0; subst; assumption.
  destruct x0; inversion Hx0.
  simpl in Hl.
  inversion Hl; subst.
  eapply IHl'; eauto.
Qed.

Definition sigma1 [A : Type] (P : A -> Prop) :=
  inhabited { Q : A -> nat -> Prop &
  ((forall (f : A) (n : nat), {Q f n} + {~Q f n})
  * (forall (f : A), P f <-> exists (n : nat), Q f n))%type }.

Definition reflectspfn [A R : Type] (P : A -> Prop) (F : A -> part R) :=
  forall (f : A), P f <-> ter (F f).

Lemma sigma1pfn [A R : Type] (P : A -> Prop) :
  inhabited R -> sigma1 P <-> exists (F : A -> part R),
    reflectspfn P F.
Proof.
  intros [r].
  split; intros HPr.
  - destruct HPr as ((Q & Hdec & Hiff)).
    pose (F f := bind (mu_tot (fun n => if Hdec f n then true else false)) (fun _ => ret r)).
    exists F.
    intros f.
    split; intros Hf.
    + rewrite Hiff in Hf.
      destruct Hf as (n & Hn).
      assert (ter (mu_tot (fun n => if Hdec f n then true else false))) as (m & Hm). {
        apply mu_tot_ter with (n := n).
        destruct (Hdec f n); auto.
      }
      exists r.
      unfold F.
      psimpl.
    + rewrite Hiff.
      destruct Hf as (r' & Hr').
      unfold F in Hr'.
      psimpl.
      exists x.
      rewrite mu_tot_hasvalue in H.
      destruct H as (H & _).
      destruct (Hdec f x); auto.
      inversion H.
  - destruct HPr as (F & HF).
    set (Q f n := if seval (F f) n then True else False).
    constructor.
    exists Q.
    split. {
      intros f n.
      unfold Q.
      destruct (seval (F f) n); auto.
    }
    intros f; split.
    + intros Hf.
      apply HF in Hf.
      destruct Hf as (r' & Hr').
      rewrite seval_hasvalue in Hr'.
      destruct Hr' as (n & Hn).
      exists n.
      unfold Q; rewrite Hn; auto.
    + intros (n & Hn).
      unfold Q in Hn.
      destruct (seval (F f) n) eqn:E; try contradiction.
      apply HF.
      exists r0.
      rewrite seval_hasvalue.
      exists n; auto.
Qed.

Definition nth_undef [A : Type] (l : list A) (n : nat) : part A :=
  match nth_error l n with
  | Some a => ret a
  | None => undef
  end.

Lemma nth_undef_app_r [A : Type] (l : list A) (n : nat) :
  n < length l -> forall (l' : list A),
    equiv (nth_undef l n) (nth_undef (l ++ l') n).
Proof.
  intros Hn l'.
  unfold nth_undef.
  rewrite nth_error_app1; eauto.
  reflexivity.
Qed.

Lemma fold_tab_nth_undef [A : Type] (l : list A) :
  forall (n : nat) (l' : list A),
    fold_part_list (tabulate (nth_undef l) 0 n) =! l' <-> n <= length l /\ l' = rev (firstn n l).
Proof.
  induction l as [|a l IHl] using rev_ind; intros n l'; cbn; eauto.
  1: destruct n as [|n]; cbn; (split; [intros | intros []]; subst; psimpl).
  destruct n as [|n].
  1:cbn; split; [intros | intros []]; subst; psimpl; eauto with arith.
  rewrite tabulate_succ', fold_part_list_app; eauto with arith.
  cbn.
  split; intros; simpl_assms.
  clear H2.
  - assert (Hn : n < length (l ++ [ a ])). {
      apply nth_error_Some.
      intros Hc.
      unfold nth_undef in H0.
      simpl in *.
      rewrite Hc in H0.
      psimpl.
    }
    split; eauto with arith.
    rewrite last_length in Hn.
    simpl in *.
    apply le_S_n in Hn.
    rewrite tabulate_equiv with (g := nth_undef l) in H.
    2: {
      intros k Hk.
      symmetry.
      apply nth_undef_app_r.
      eauto with arith.
    }
    apply IHl in H as (_ & ->).
    erewrite list_basics.take_S_r.
    1:rewrite <- firstn_app_l, rev_unit; eauto.
    destruct (list_basics.nth_lookup_or_length (l ++ [a]) n x1) as [H | H].
    2: rewrite last_length in H; lia.
    unfold nth_undef in H0.
    destruct (nth_error _ _) eqn:E; psimpl.
    erewrite nth_error_nth in H; eauto.
  - destruct H as (Hn & Hl').
    rewrite last_length in Hn.
    apply le_S_n in Hn.
    destruct (base.lookup n (l ++ [a])) as [a'|] eqn:E.
    2: {
      apply list_basics.lookup_ge_None in E.
      rewrite last_length in E.
      lia.
    }
    erewrite list_basics.take_S_r, rev_unit in Hl'; eauto.
    rewrite <- firstn_app_l in Hl'; eauto.
    assert (Hl : fold_part_list (tabulate (nth_undef l) 0 n) =! rev (firstn n l)) by (apply IHl; eauto).
    rewrite tabulate_equiv with (g := nth_undef l).
    2: {
      intros k Hk.
      symmetry.
      apply nth_undef_app_r.
      eauto with arith.
    }
    assert (nth_undef (l ++ [a]) n =! a'). {
      unfold nth_undef.
      destruct (nth_error _ _) eqn:E'.
      - apply nth_error_nth with (d := a) in E'.
        apply list_basics.nth_lookup_Some with (d := a) in E.
        subst; psimpl.
      - apply nth_error_None in E'.
        rewrite last_length in E'; lia.
    }
    subst.
    psimpl.
Qed.

End PartialZoo.