Require Import ssreflect Setoid Lia List.
Require Import PostTheorem.external.Shared.embed_nat.
From PostTheorem Require Import external.mu_nat.
From Coq.Logic Require Import ConstructiveEpsilon.
Import EmbedNatNotations ListNotations.

Local Set Implicit Arguments.
Local Unset Strict Implicit.

Definition monotonic {X} (f : option X) :=
   n1 x, f = Some x n2, f = Some x.

Definition agnostic {X} (f : option X) :=
   n1 n2 y1 y2, f = Some f = Some = .

Lemma monotonic_agnostic {X} (f : option X) :
  monotonic f agnostic f.
Proof.
  move mono .
  destruct (Compare_dec.le_ge_dec ) as [l | g].
  - eapply mono in l; eauto. congruence.
  - eapply mono in g; eauto. congruence.
Qed.


Class partiality :=
  {

    part : Type Type ;

    hasvalue : A, part A A Prop ;
    hasvalue_det : A x (a1 a2 : A), hasvalue x hasvalue x = ;

    ret : A, A part A ;
    ret_hasvalue : A (a : A), hasvalue (ret a) a ;

    bind : A B, part A (A part B) part B ;
    bind_hasvalue : A B x(f : A part B) b, hasvalue (bind x f) b a, hasvalue x a hasvalue (f a) b;

    undef : {A}, part A ;
    undef_hasvalue : A (a : A), hasvalue undef a ;

    mu : ( part bool) part ;
    mu_hasvalue : (f : part bool) n, hasvalue (mu f) n (hasvalue (f n) true m, m < n hasvalue (f m) false);
  
    seval : A, part A option A ;
    seval_mono : A x, monotonic (@seval A x) ;
    seval_hasvalue : A x (a : A), hasvalue x a n, seval x n = Some a ;

  }.

Definition ter {Part : partiality} A (x : part A) := a, hasvalue x a.
Definition equiv {Part : partiality} A (x y : part A) := a, hasvalue x a hasvalue y a.

Notation "a =! b" := (hasvalue a b) (at level 50).

Notation "A ↛ B" := (A part B) (at level 10).

Definition pcomputes {X Y} `{partiality} (f : X Y) (R : X Y Prop) :=
   x y, f x =! y R x y.

Definition functional {X Y} (R : X Y Prop) :=
   x y1 y2, R x R x = .

Definition total {X Y} (R : X Y Prop) :=
   x, y, R x y.

Record FunRel X Y := {the_rel :> X Y Prop ; the_func_proof : functional the_rel}.
Arguments the_rel {_ _}.

Global Instance part_equiv_Equivalence `{partiality} {A} :
  Equivalence (@equiv _ A).
Proof.
  firstorder.
Qed.


Section assume_part.

  Context {impl : partiality}.

  Lemma ret_hasvalue_inv {A} (a1 a2 : A) :
    ret =! = .
  Proof.
    move H.
    eapply hasvalue_det. eapply ret_hasvalue. eauto.
  Qed.


  Lemma ret_hasvalue' {A} (a1 a2 : A) :
     = ret =! .
  Proof.
    intros . eapply ret_hasvalue.
  Qed.


  Lemma ret_hasvalue_iff {A} (a1 a2 : A) :
     = ret =! .
  Proof.
    split.
    - apply ret_hasvalue'.
    - apply ret_hasvalue_inv.
  Qed.


  Definition mu_tot (f : bool) := mu ( n ret (f n)).

  Lemma mu_tot_hasvalue (f : bool) n :
    hasvalue (mu_tot f) n (f n = true m, m < n f m = false).
  Proof.
    unfold mu_tot. rewrite mu_hasvalue. now repeat setoid_rewrite ret_hasvalue_iff.
  Qed.


  Lemma mu_tot_ter (f : bool) n :
    f n = true
    ter (mu_tot f).
  Proof.
    move H.
    assert (He : n, f n = true) by eauto.
    assert (d : n, {f n = true} + { f n = true}) by (move n'; destruct (f n'); firstorder congruence).
    destruct (mu_nat_dep _ d He) as [n' Hn'] eqn:E.
    eapply (f_equal (@proj1_sig _ _)) in E.
    exists n'. eapply mu_tot_hasvalue. split.
    - eauto.
    - move m Hlt. cbn in E. subst.
      eapply mu_nat_dep_min in Hlt. destruct (f m); congruence.
  Qed.


  Definition undef' : {A}, A part A := A a0 bind (mu_tot ( _ false)) ( n ret ).

  Lemma undef'_hasvalue : A a0 (a : A), hasvalue (undef' ) a.
  Proof.
    intros A a [a' [[[=]] % mu_tot_hasvalue ]] % bind_hasvalue.
  Qed.


  Definition eval' {A} (x : part A) (H : ter x) : {a : A | hasvalue x a}.
  Proof.
    assert (Hx : n, seval x n None). {
      destruct H as [a [n H] % seval_hasvalue]. exists n. congruence.
    }
    eapply constructive_indefinite_ground_description_nat in Hx as [n Hx].
    - destruct seval eqn:E; try congruence. exists a. eapply seval_hasvalue. firstorder.
    - move n. destruct seval; firstorder congruence.
  Qed.


  Definition eval {A} (x : part A) (H : ter x) : A := proj1_sig (eval' H).
  Definition eval_hasvalue {A} (x : part A) (H : ter x) : hasvalue x (eval H) := proj2_sig (eval' H).

  Definition mkpart {A} (f : option A) :=
    bind (mu_tot ( n match f n with Some a true | None false end))
      ( n match f n with Some a ret a | None undef end).

  Lemma mkpart_hasvalue1 {A} (f : option A) :
     a, mkpart f =! a n, f n = Some a.
  Proof.
    move a.
    rewrite /mkpart bind_hasvalue.
    move [] x [] / mu_hasvalue [] ter_mu Hmu Hf.
    exists x. destruct (f x). eapply (hasvalue_det (ret_hasvalue )) in Hf as .
    reflexivity. eapply undef_hasvalue in Hf as [].
  Qed.


  Lemma mkpart_ter {A} (f : option A) n a :
    f n = Some a ter (mkpart f).
  Proof.
    move Hn. unfold ter.
    rewrite /mkpart. setoid_rewrite bind_hasvalue.
    assert (Hf : n, f n None). exists n. firstorder congruence.
    assert (d : n : , {( n0 : f None) n} + { ( n0 : f None) n}). { move . destruct (f ); firstorder congruence. }
    edestruct (mu_nat_dep _ d Hf) as [m Hm] eqn:E. eapply (f_equal (@proj1_sig _ _)) in E. cbn in E.
    destruct (f m) as [|]eqn:; try congruence.
    exists , m.
    rewrite mu_tot_hasvalue. repeat split.
    + rewrite . reflexivity.
    + subst. move m' Hle.
      destruct (f m') eqn:.
      * eapply mu_nat_dep_min in Hle. firstorder congruence.
      * reflexivity.
    + rewrite . eapply ret_hasvalue.
  Qed.


  Lemma mkpart_hasvalue2 {A} (f : option A) n a :
    agnostic f f n = Some a mkpart f =! a.
  Proof.
    move Hagn Hn.
    destruct (mkpart_ter Hn) as [a' H].
    destruct (mkpart_hasvalue1 H) as [n' H'].
    now assert (a' = a) as by (eapply Hagn; eauto).
  Qed.


  Lemma mkpart_hasvalue {A} (f : option A) :
    agnostic f a, mkpart f =! a n, f n = Some a.
  Proof.
    split.
    eapply mkpart_hasvalue1.
    move [n Hn]. eapply mkpart_hasvalue2; eauto.
  Qed.


  Definition par : A B, part A part B part (A + B) :=
     A B x y
    bind (mu_tot ( n if seval x n is Some a then true else if seval y n is Some b then true else false))
      ( n if seval x n is Some a then ret (inl a) else if seval y n is Some b then ret (inr b) else undef).

  Lemma par_hasvalue1 : A B (x : part A) (y : part B) a, hasvalue (par x y) (inl a) hasvalue x a.
  Proof.
    intros A B x y a [a' [( & ) % mu_tot_hasvalue ]] % bind_hasvalue.
    destruct (seval x a') eqn:, (seval y a') eqn:; try congruence.
    - eapply ret_hasvalue_iff in as [= ]. eapply seval_hasvalue. eauto.
    - eapply ret_hasvalue_iff in as [= ]. eapply seval_hasvalue. eauto.
    - eapply ret_hasvalue_iff in as [= ].
  Qed.


  Lemma par_hasvalue2 : A B (x : part A) (y : part B) b, hasvalue (par x y) (inr b) hasvalue y b.
  Proof.
    intros A B x y b [a' [( & ) % mu_tot_hasvalue ]] % bind_hasvalue.
    destruct (seval x a') eqn:, (seval y a') eqn:; try congruence.
    - eapply ret_hasvalue_iff in as [= ].
    - eapply ret_hasvalue_iff in as [= ].
    - eapply ret_hasvalue_iff in as [= ]. eapply seval_hasvalue. eauto.
  Qed.


  Lemma par_hasvalue3 : A B (x : part A) (y : part B), ter x ter y ter (par x y).
  Proof.
    intros A B x y [[a H] | [b H]].
    - eapply seval_hasvalue in H as [n Hn].
      destruct (@mu_tot_ter ( n if seval x n is Some a then true else if seval y n is Some b then true else false) n) as [m Hm].
      + now rewrite Hn.
      + pose proof (Hm' := Hm).
        eapply mu_tot_hasvalue in Hm as ( & ).
        destruct (seval x m) eqn:, (seval y m) eqn:; try congruence.
        * exists (inl ). eapply bind_hasvalue. eexists. split. eapply Hm'. rewrite . eapply ret_hasvalue.
        * exists (inl ). eapply bind_hasvalue. eexists. split. eapply Hm'. rewrite . eapply ret_hasvalue.
        * exists (inr b). eapply bind_hasvalue. eexists. split. eapply Hm'. rewrite . eapply ret_hasvalue.
    - eapply seval_hasvalue in H as [n Hn].
      destruct (@mu_tot_ter ( n if seval x n is Some a then true else if seval y n is Some b then true else false) n) as [m Hm].
      + rewrite Hn. clear. now destruct seval.
      + pose proof (Hm' := Hm).
        eapply mu_tot_hasvalue in Hm as ( & ).
        destruct (seval x m) eqn:, (seval y m) eqn:; try congruence.
        * exists (inl a). eapply bind_hasvalue. eexists. split. eapply Hm'. rewrite . eapply ret_hasvalue.
        * exists (inl a). eapply bind_hasvalue. eexists. split. eapply Hm'. rewrite . eapply ret_hasvalue.
        * exists (inr b). eapply bind_hasvalue. eexists. split. eapply Hm'. rewrite .
          erewrite (monotonic_agnostic (@seval_mono _ _ _) Hn ). eapply ret_hasvalue.
  Qed.


End assume_part.