From Coq Require Import List Morphisms Recdef ZArith Znat Arith.
From QuickChick  Require Import Sets Tactics Producer Generators Enumerators.

Set Bullet Behavior "Strict Subproofs".

Local Open Scope set_scope.

(** Apply a function n times *)
Fixpoint appn {A} (f : A -> A) (n : nat) : A ->  A :=
  fun x =>
    match n with
      | 0%nat => x
      | S n' => f (appn f n' x)
    end.

Infix "^" := appn (at level 30, right associativity) : fun_scope.


(** Instance Hierarchy  

   GenSized 
      |
      |
     Gen   Shrink
       \    /
        \  /
      Arbitrary
*)

(** * Generator-related classes *)

(* begin gen_sized_class *)
Class GenSized (A : Type) :=
  { arbitrarySized : nat -> G A }.
(* end gen_sized_class *)

(* begin gen_class *)
Class Gen (A : Type) := { arbitrary : G A }.
(* end gen_class *)

(** Shrink class *)
Class Shrink (A : Type) :=
  { shrink : A -> list A }.

(** Arbitrary Class *)
Class Arbitrary (A : Type) `{Gen A} `{Shrink A}.

Class EnumSized (A : Type) := { enumSized : nat -> E A }.
  
Class Enum (A : Type) := { enum : E A }.

(* ZP: This is not longer usefull *) 
(* 
(** * Sizes of types *)
  
Class Sized (A : Type) :=
  { size : A -> nat }.

Class CanonicalSized (A : Type) `{Sized A} :=
  {
    zeroSized : set A;
    succSized : set A -> set A;

    zeroSized_spec : zeroSized <--> [ set x : A | size x = 0 ];
    succSized_spec :
      forall n, succSized [ set x : A | size x <= n ] <--> [ set x : A | size x <= S n ]
 
  }.

Lemma size_ind (A : Type) `{Hyp : Sized A} :
  forall (P : A -> Prop), (forall y, (forall x, size x < size y -> P x) -> P y) -> (forall x, P x).
Proof.
  intros P H1.
  intros x.
  assert (Hin : [ set y :  A | size y <= size x] x); eauto.
  revert Hin.
  generalize (size x). intros n.
  revert x. induction n.
  - intros x Hl. apply H1. intros x1 Hlt. ssromega.
  - intros x Hleq. eapply H1. intros x1 Hlt.
    eapply IHn. ssromega.
Qed.

Lemma size_lfp (A : Type) `{Hyp : CanonicalSized A} :
  [set x : A | True ] <--> \bigcup_(s : nat) [set x : A | size x <= s ].
Proof.
  intros a; split; eauto. intros _.
  exists (size a). split; eauto. constructor.
Qed.

Lemma succ_lfp (A : Type) `{Hyp : CanonicalSized A}
      `{Proper _ (respectful set_eq set_eq) succSized} s :
  [set x : A | size x <= s ] <-->  (succSized ^ s) zeroSized.
Proof.
  induction s.
  simpl.
  - rewrite zeroSized_spec.
    split; intros; ssromega.
  - simpl. rewrite <- succSized_spec.
    rewrite IHs. reflexivity.
Qed.
*) 
(* Lemma succ_lfp' (A : Type) `{Hyp : CanonicalSized A} : *)
(*   \bigcup_(s : nat)  (succSized ^ s) zeroSized <--> [ set x : A | True ]. *)
(* Proof. *)
(*   intros. split; eauto. *)
(*   intros _. *)
(*   eapply set_eq_trans. *)
(*   Focus 2. symmetry. *)
(*   eapply succ_lfp. *)
(*   simpl.  *)
(*   rewrite succ_lfp at 2. *)
(* split. *)
(*     split. rewrite IHs. firstorder. *)
(*     IHs. *)
(*   firstorder. reflexivity. split; intros; eauto. *)
(*   exists (size a). *)
(*   remember (size a) as s. *)
(*   revert a Heqs. induction s; intros. *)
(*   - split. constructor. *)
(*     simpl. eapply zeroSized_spec. now eauto. *)
(*   - split. constructor. *)
(*     simpl. *)
(*     eapply (succSized_spec. *)
(*     eassumption. *)
(*   eapply size_ind. *)
  
(*   [set x : A | True ] <-->  . *)

(** * Correctness classes *)

(** Correctness of sized generators *)

Class CorrectSized {A : Type} {G} `{Producer G} (g : nat -> G A)  :=
  { prodCorrectSized :
      [ set n | exists s, semProd (g s) n ] <--> [set : A]
  }.


(** Correctness of generators *)
Class Correct (A : Type) {G} `{Producer G} (g : G A)  :=
  {
    prodCorrect : semProd g <--> [set : A]
  }.

(** * Monotonic generators *)

(** Monotonicity of size parametric generators *)
Class GenSizedMonotonic (A : Type) `{GenSized A}
      `{forall s, SizeMonotonic (arbitrarySized s)}.

(** Monotonicity of size parametric generators v2 *)
Class GenSizedSizeMonotonic (A : Type) `{GenSized A}
        `{SizedMonotonic A G arbitrarySized}.

Class GenMonotonic (A : Type) `{Gen A} `{@SizeMonotonic A G ProducerGen arbitrary}.

(** Monotonicity of size parametric generators *)
Class EnumSizedMonotonic (A : Type) `{EnumSized A}
      `{forall s, @SizeMonotonic A E ProducerEnum (enumSized s)}.

(** Monotonicity of size parametric generators v2 *)
Class EnumSizedSizeMonotonic (A : Type) `{EnumSized A}
        `{@SizedMonotonic A E ProducerEnum enumSized}.

Class EnumMonotonic (A : Type) `{Enum A}
        `{@SizeMonotonic A E ProducerEnum enum}.
  
(** * Correct generators *)

Class GenSizedCorrect (A : Type) `{GenSized A}
      `{@CorrectSized A G ProducerGen arbitrarySized}.

Class GenCorrect (A : Type) `{Gen A}
        `{@Correct A G ProducerGen arbitrary}.
 
(* Monotonic and Correct generators *)
Class GenMonotonicCorrect (A : Type)
      `{Gen A}
      `{@SizeMonotonic A G ProducerGen arbitrary}
      `{@Correct  A G ProducerGen arbitrary}.

Class EnumSizedCorrect (A : Type)
        `{EnumSized A} `{@CorrectSized A E ProducerEnum enumSized}.

Class EnumCorrect (A : Type) `{Enum A}
        `{@Correct A E ProducerEnum enum}.
 
(* Monotonic and Correct generators *)
Class EnumMonotonicCorrect (A : Type)
      `{Enum A}
      `{@SizeMonotonic A E ProducerEnum enum}
      `{@Correct A E ProducerEnum enum}.

  
(** Coercions *)
  
#[global]
Instance GenSizedMonotonicOfSizeMonotonic
         (A : Type) (Hgen : GenSized A) (Hmon : forall s, @SizeMonotonic A G ProducerGen (arbitrarySized s))
: @GenSizedMonotonic A Hgen Hmon := {}.
  
#[global]
Instance GenMonotonicOfSizeMonotonic
         (A : Type) (Hgen : Gen A) (Hmon : @SizeMonotonic A G ProducerGen arbitrary)
: @GenMonotonic A Hgen Hmon := {}.

#[global]
Instance GenSizedCorrectOfSizedCorrect
         (A : Type) (Hgen : GenSized A)
         `{Hcor : @CorrectSized A G ProducerGen arbitrarySized}
: @GenSizedCorrect A Hgen Hcor := {}.

#[global]
Instance GenMonotonicCorrectOfMonotnicCorrect
         (A : Type) (Hgen : Gen A)
         (Hmon : @SizeMonotonic A G ProducerGen arbitrary)
         `{Hcor : @Correct A G ProducerGen arbitrary}
  : @GenMonotonicCorrect A Hgen Hmon Hcor := {}.



#[global]
Instance GenSizedSizeMonotonicOfSizedMonotonic
         (A : Type) (Hgen : GenSized A) (Hmon : @SizedMonotonic A G ProducerGen arbitrarySized)
: @GenSizedSizeMonotonic A Hgen ProducerGen Hmon := {}.

#[global] Instance GenOfGenSized {A} `{GenSized A} : Gen A :=
  {| arbitrary := sized arbitrarySized |}.

#[global] Instance ArbitraryOfGenShrink {A} `{Gen A} `{Shrink A} : Arbitrary A := {}.

Generalizable Variables PSized PMon PSMon PCorr.

#[global] Instance GenMonotonicOfSized (A : Type)
         `{H : GenSized A}
         `{@GenSizedMonotonic A H PMon}
         `{@GenSizedSizeMonotonic A H ProducerGen PSMon}
  : @GenMonotonic A
                  (@GenOfGenSized A H) 
                  (@sizedSizeMonotonic G ProducerGen _ A
                                       (@arbitrarySized A H)
                                       PMon PSMon) := {}.

#[global]
Instance GenCorrectOfSized (A : Type)
         {H : GenSized A}
         `{@GenSizedMonotonic A H PMon}
         `{@GenSizedSizeMonotonic A H ProducerGen PSMon}
         `{@GenSizedCorrect A H PCorr} : Correct A arbitrary.
Proof.
  constructor. unfold arbitrary, GenOfGenSized. 
  eapply set_eq_trans.
  - eapply semSized_alt; now eauto with typeclass_instances.
  - destruct PCorr. 
    rewrite <- prodCorrectSized0. 

    split. intros [n [? H3]]. eexists. eassumption.

    intros [s H4]. eexists; split; eauto. reflexivity.
Qed.

(* 
Lemma nat_set_ind (A : Type) `{GenSized A} `{Hyp : CanonicalSized A} :
  (semProd (arbitrarySized 0) <--> zeroSized) ->
  (forall (s : nat) (elems : set A),
     semProd (arbitrarySized s) <--> elems ->
     semProd (arbitrarySized (s.+1)) <--> succSized elems) ->
  (forall s : nat, semProd (arbitrarySized s) <--> (fun x : A => size x <= s)).
Proof.
  intros HO IH. intros n; induction n.
  - eapply set_eq_trans with (B := (fun x : A => size x = 0)).
    rewrite -zeroSized_spec //=.
    intros s. destruct (size s). now firstorder.
    split; intros; ssromega.
  - rewrite -succSized_spec. eauto.
Qed.
 *)

#[global] Instance EnumSizedMonotonicOfSizeMonotonic
         (A : Type) (Hgen : EnumSized A) (Hmon : forall s, @SizeMonotonic A E ProducerEnum (enumSized s))
: @EnumSizedMonotonic A Hgen Hmon := {}.
  
#[global] Instance EnumMonotonicOfSizeMonotonic
         (A : Type) (Hgen : Enum A) (Hmon : @SizeMonotonic A E ProducerEnum enum)
: @EnumMonotonic A Hgen Hmon := {}.

#[global] Instance EnumSizedCorrectOfSizedCorrect
         (A : Type) (Hgen : EnumSized A)
         `{Hcor : @CorrectSized A E ProducerEnum enumSized}
: @EnumSizedCorrect A Hgen Hcor := {}.

#[global] Instance EnumCorrectOfCorrect
         (A : Type) (Hgen : Enum A)
         `{Hcor : @Correct A E ProducerEnum enum}
: @EnumCorrect A Hgen Hcor := {}.

#[global] Instance EnumMonotonicCorrectOfMonotnicCorrect
         (A : Type) (Hgen : Enum A)
         (Hmon : @SizeMonotonic A E ProducerEnum enum)
         `{Hcor : @Correct A E ProducerEnum enum}
  : @EnumMonotonicCorrect A Hgen Hmon Hcor := {}.


#[global] Instance EnumSizedSizeMonotonicOfSizedMonotonic
         (A : Type) (Hgen : EnumSized A) (Hmon : @SizedMonotonic A E ProducerEnum enumSized)
: @EnumSizedSizeMonotonic A Hgen Hmon := {}.

#[global] Instance EnumOfEnumSized {A} `{EnumSized A} : Enum A :=
  {| enum := sized enumSized |}.

(*
#[global] Instance EnumOfGenShrink {A} `{Gen A} `{Shrink A} : Arbitrary A := {}.
 *)

#[global] Instance EnumMonotonicOfSized (A : Type)
         `{H : EnumSized A}
         `{@EnumSizedMonotonic A H PMon}
         `{@EnumSizedSizeMonotonic A H PSMon}
  : @EnumMonotonic A
                  (@EnumOfEnumSized A H)
                  (@sizedSizeMonotonic E ProducerEnum _ A
                                       (@enumSized A H)
                                       PMon PSMon) := {}.

#[global] Instance EnumCorrectOfSized' (A : Type)
         {H : EnumSized A}
         `{@EnumSizedMonotonic A H PMon}
         `{@EnumSizedSizeMonotonic A H PSMon}
         `{@EnumSizedCorrect A H PCorr} : Correct A enum.
Proof.
  constructor. unfold arbitrary, EnumOfEnumSized. 
  eapply set_eq_trans.
  - eapply semSized_alt; eauto with typeclass_instances.
  - destruct PCorr. 
    rewrite <- prodCorrectSized0. 

    split. intros [n [? H3]]. eexists. eassumption.

    intros [s H4]. eexists; split; eauto. reflexivity.
Qed.

#[global] Instance EnumCorrectOfSized (A : Type)
         (H1 : EnumSized A)
         (H3 : forall s : nat, SizeMonotonic (enumSized s))
         (H4 : SizedMonotonic enumSized)
         (H5 : CorrectSized enumSized)
  : Correct A enum.
Proof.
  eapply EnumCorrectOfSized'; eauto.
  constructor; eauto.
  constructor; eauto.
  constructor; eauto.
Qed.


Lemma enumerating_complete A (e : E A) {Hc : Correct A e} ch :
  (exists x, ch x = Some true) ->
  exists s, enumerating e ch s = Some true.
Proof.
  intros [x Hch].
  unfold enumerating.
  assert (Hin : semProd e x).
  { eapply Hc. reflexivity. }
  destruct Hin as [s [_ Hsem]]. simpl in *.
  unfold semEnumSize in *.
  exists s. revert Hsem.
  generalize (Enumerators.run e s), false.
  induction l; intros b Hin; inv Hin; simpl.
  - rewrite Hch. reflexivity.
  - destruct (ch a); eauto.
    destruct b0; eauto.
Qed.

Lemma lazylist_backtrack_opt_is_true A (l: LazyList.LazyList (option A)) ch b :
  lazylist_backtrack_opt l ch true = Some b -> b = true. 
Proof.
  induction l.
  - simpl. congruence.
  - simpl in *.
    destruct a.
    + destruct (ch a) eqn:Heq. destruct b0. congruence.
      now eauto. now eauto.
    + eauto.
Qed. 

Lemma lazylist_backtrack_opt_is_false A (l: LazyList.LazyList (option A)) ch b :
  lazylist_backtrack_opt l ch b = Some false -> b = false. 
Proof.
  revert b; induction l; intros b.
  - simpl. destruct b; try congruence.
  - simpl in *. 
    destruct a.
    + destruct (ch a) eqn:Heq. destruct b0. congruence.
      now eauto. intros Hin. eapply H in Hin. congruence. 
    + intros Hn. eapply lazylist_backtrack_opt_is_true in Hn.
      congruence.
Qed. 

Lemma lazylist_backtrack_opt_true A (l: LazyList.LazyList (option A)) ch b :
  lazylist_backtrack_opt l ch b = Some true <->
  exists x, LazyList.In_ll (Some x) l /\ ch x = Some true.
Proof.
  revert b; induction l; intros b.
  - simpl. split; intros Hc; try congruence.
    destruct b; congruence. 
    destruct Hc. inv H; exfalso; eauto.
  - split.
    + intros Hl.
      simpl in *. destruct a.
      destruct (ch a) eqn:Heq.
      * destruct b0.
        -- eexists. split; eauto.
        -- eapply H in Hl. destruct Hl. inv H0.
           eexists; split; eauto.
      * eapply H in Hl. destruct Hl. inv H0.
        eexists; split; eauto.
      * eapply H in Hl. destruct Hl. inv H0.
        eexists; split; eauto.
    + intros Hc; inv Hc. inv H0. simpl in *. destruct a. 
      * inv H1. inv H3. rewrite H2. reflexivity.
        destruct (ch a) eqn:Heq. destruct b0. reflexivity.
        eapply H. eexists. now split; eauto.
        eapply H. eexists. now split; eauto.
      * inv H1. congruence. eapply H. eexists. now split; eauto.
Qed.

Lemma lazylist_backtrack_opt_false A (l: LazyList.LazyList (option A)) ch :
  lazylist_backtrack_opt l ch false = Some false <->
  (forall x, LazyList.In_ll (Some x) l -> ch x = Some false) /\
  ~ LazyList.In_ll None l.
Proof.
  induction l. 
  - simpl. split; intros Hc.
    split; eauto. intros. now inv H.
    reflexivity.
  - split.
    + intros Hl.
      simpl in *.
      destruct a.
      * destruct (ch a) eqn:Heq. destruct b.
        congruence.
        
        -- eapply H in Hl. inv Hl.
           split.
           intros. inv H2; eauto. inv H3. eassumption.

           intros Hc; inv Hc. congruence. eauto.

        -- eapply lazylist_backtrack_opt_is_false in Hl. congruence.

      * eapply lazylist_backtrack_opt_is_false in Hl. congruence.

    + intros [H1 H2].
      simpl.
      destruct a. 
      * destruct (ch a) eqn:Heq. destruct b.
        rewrite H1 in Heq. congruence.
        now left.
        eapply H. split. intros. eapply H1.
        right. eassumption.
        intros Hc. eapply H2. right. eassumption.
        rewrite H1 in Heq. congruence. now left.

      * exfalso. eapply H2. now left.
Qed.        
        

Lemma enumeratingOpt_monotonic A (e : E (option A))
      {Hc : SizeMonotonicOpt e} {Hfp : SizeFP e}
      ch1 ch2 s1 s2 b :
  s1 <= s2 ->
  (forall x, ch1 x = Some b -> ch2 x = Some b) ->
  enumeratingOpt e ch1 s1 = Some b ->
  enumeratingOpt e ch2 s2 = Some b.
Proof.
  intros Hleq Hall Hen.

  specialize (Hc _ _ Hleq).
  specialize (Hfp _ _ Hleq).
  
  unfold enumeratingOpt in *.
  unfold semProdSizeOpt in *. simpl in *. 
  unfold semEnumSize in *.
  revert Hc Hfp Hen. 

  generalize (run e s1), (run e s2).

  intros l1 l2 Hs1 Hs2 Hl.
  destruct b.

  - eapply lazylist_backtrack_opt_true in Hl.
    destruct Hl. destruct H.

    eapply lazylist_backtrack_opt_true.
    eexists. split.
    eapply Hs1. eassumption.
    eapply Hall. eassumption.

  - eapply lazylist_backtrack_opt_false in Hl.
    destruct Hl. 

    eapply lazylist_backtrack_opt_false.

    split. intros x Hin. eapply Hall. eapply H.
    eapply Hs2. eassumption. eassumption.

    intros Hc. eapply H0. eapply Hs2. eassumption. eassumption.
Qed.
