Set Warnings "-notation-overridden,-parsing".

From Coq Require Import String List ZArith Lia ssreflect ssrfun ssrbool.
From mathcomp Require Import ssrnat eqtype seq.
Set SsrOldRewriteGoalsOrder.  (* change Set to Unset when porting the file, then remove the line when requiring MathComp >= 2.6 *)

From QuickChick Require Import
     RandomQC RoseTrees Sets Tactics Producer.

Set Implicit Arguments.
Unset Strict Implicit.
Unset Printing Implicit Defensive.

Set Bullet Behavior "Strict Subproofs".

Import ListNotations.

(* Low-level Generators *)

Local Open Scope set_scope.

Inductive GenType (A:Type) : Type := MkGen : (nat -> RandomSeed -> A) -> GenType A.
  
Definition G := GenType.

(** * Primitive generator combinators *)
  
Definition run {A : Type} (g : G A) := match g with MkGen f => f end.
  
Definition returnGen {A : Type} (x : A) : G A :=
  MkGen (fun _ _ => x).

Definition bindGen {A B : Type} (g : G A) (k : A -> G B) : G B :=
  MkGen (fun n r =>
             let (r1,r2) := randomSplit r in
             run (k (run g n r1)) n r2).

#[global] Instance MonadGen : Monad G :=
  { ret := @returnGen
  ; bind := @bindGen }.

Fixpoint createRange (n : nat) (acc : list nat) : list nat :=
  match n with
  | O => List.rev (cons O acc)
  | S n' => createRange n' (cons n acc)
  end.

Fixpoint rnds (s : RandomSeed) (n' : nat) : list RandomSeed :=
    match n' with
      | O => nil
      | S n'' =>
        let (s1, s2) := randomSplit s in
        cons s1 (rnds s2 n'')
    end.

Definition sampleGen (A : Type) (g : G A) : list A :=
    match g with
      | MkGen m =>
        let rnd := newRandomSeed in
        let l := List.combine (rnds rnd 20) (createRange 10 nil) in
        List.map (fun (p : RandomSeed * nat) => let (r,n) := p in m n r) l
    end.

Definition sizedGen {A : Type} (f : nat -> G A) : G A :=
  MkGen (fun n r => run (f n) n r).

Definition resizeGen {A : Type} (n : nat) (g : G A) : G A :=
    match g with
    | MkGen m => MkGen (fun _ => m n)
    end.

Definition semGenSize {A : Type} (g : G A) (s : nat) : set A := codom (run g s).

Definition chooseGen {A : Type} {le} `{ChoosableFromInterval A le} (range : A * A) : G A :=
    MkGen (fun _ r => fst (randomR range r)).

#[global] Program Instance ProducerGen : Producer G :=
  {
  super := MonadGen;

  sample := sampleGen;
  
  sized  := @sizedGen; 
  resize := @resizeGen;

  choose := @chooseGen;
  
  semProdSize := @semGenSize;

  (* Probably belongs in another class for modularity? *)
  bindPf :=
    fun {A B : Type} (g : G A)
        (k : forall (a : A),
            (fun (A : Type) (g : G A) =>
               \bigcup_(size in [set: nat]) semGenSize g size) A g a -> G B) =>
    MkGen (fun n r =>
             let (r1,r2) := randomSplit r in
             run (k (run g n r1) _) n r2)}.
Next Obligation.
  unfold semGenSize, codom, bigcup.
  exists n; split => //=.
  exists r1; auto.
Defined.

(* Generator specific sample of a single input. *)
Definition sample1 (A : Type) (g : G A) : A :=
    match g with
      | MkGen m =>
        let rnd := newRandomSeed in
        m 10 rnd
    end.

Lemma runFmap (A B : Type) (f : A -> B) (g : G A) seed size :
    run (fmap f g) size seed = f (run g size (fst (randomSplit seed))).
Proof. simpl.
       destruct (randomSplit seed).
       simpl.
       auto.
Qed.

(* Generator specific - shrinking support. *)
Definition promote {A : Type} (m : Rose (G A)) : G (Rose A) :=
    MkGen (fun n r => fmapRose (fun g => run g n r) m).

(* Generator specific - coarbitrary support. *)
Definition variant {A : Type} (p : SplitPath) (g : G A) : G A := 
    match g with 
      | MkGen f => MkGen (fun n r => f n (varySeed p r))
    end.
  
Definition reallyUnsafeDelay {A : Type} : G (G A -> A) :=
    MkGen (fun r n g => (match g with MkGen f => f r n end)).
  
Definition reallyUnsafePromote {r A : Type} (m : r -> G A) : G (r -> A) :=
    (bindGen reallyUnsafeDelay (fun eval => 
                                  returnGen (fun r => eval (m r)))).

Lemma promoteVariant :
    forall {A B : Type} (a : A) (f : A -> SplitPath) (g : G B) size
      (r r1 r2 : RandomSeed),
      randomSplit r = (r1, r2) ->
      run (reallyUnsafePromote (fun a => variant (f a) g)) size r a =
      run g size (varySeed (f a) r1).
Proof.
    move => A B a p g size r r1 r2 H.
    rewrite /reallyUnsafePromote /variant.
    destruct g. rewrite /= H. by [].
Qed.

Lemma semPromote A (m : Rose (G A)) :
    semProd (promote m) <-->
    codom2 (fun size seed => fmapRose (fun g => run g size seed) m).
Proof. by rewrite /codom2 curry_codom2l. Qed.

Lemma semPromoteSize (A : Type) (m : Rose (G A)) n :
    semProdSize (promote m) n <-->
               codom (fun seed => fmapRose (fun g => run g n seed) m).
Proof. by []. Qed.

Lemma runPromote A (m : Rose (G A)) seed size :
    run (promote m) seed size = fmapRose (fun (g : G A) => run g seed size) m.
Proof. by []. Qed.

(* Generator specific - choose and its semantics. *)
Lemma semChooseSizeGen A {le} `{ChoosableFromInterval A le} (a1 a2 : A) :
    le a1 a2 ->
    forall size, semProdSize (choose (a1,a2)) size <-->
                       [set a | le a1 a /\ le a a2].
Proof. by move=> /= le_a1a2 m n; rewrite (randomRCorrect n a1 a2). Qed.
  
#[global] Instance chooseUnsized {A} {le} `{RandomQC.ChoosableFromInterval A le} (a1 a2 : A) :
    Unsized (choose (a1, a2)).
Proof. by []. Qed.
  
Lemma semChooseGen A {le} `{RandomQC.ChoosableFromInterval A le} (a1 a2 : A) :
    le a1 a2 ->
    (semProd (choose (a1,a2)) <--> [set a | le a1 a /\ le a a2]).
Proof.
  move=> /= le_a1a2. rewrite <- (unsized_alt_def 1).
  move => m /=. rewrite (randomRCorrect m a1 a2) //.
Qed.

Lemma semChooseSizeGenNat (a1 a2 : nat) : a1 <= a2 -> forall size,
    (semProdSize (choose (a1,a2)) size <--> [set a | a1 <= a <= a2]).
Proof.
  move => /leP H0 size. rewrite (semChooseSizeGen (H := ChooseNat) H0). intros a; apply lele_coq_ssr.
Qed.

Lemma semChooseGenNat (a1 a2 : nat) : a1 <= a2 ->
    (semProd (choose (a1,a2)) <--> [set a | a1 <= a <= a2]).
Proof.
  move => /leP H0. rewrite (semChooseGen (H := ChooseNat) H0). intros a; apply lele_coq_ssr.
Qed.

  Definition thunkGen {A} (f : unit -> G A) : G A :=
    MkGen (fun n r => run (f tt) n r).

  Lemma semThunkGenSize {A} (f : unit -> G A) s :
    semProdSize (thunkGen f) s <--> semProdSize (f tt) s.
  Proof.
    split; intros [r Hr]; exists r; simpl in *; assumption.
  Qed.

  Lemma semThunkGen {A} (f : unit -> G A) :
    semProd (thunkGen f) <--> semProd (f tt).
  Proof.
    split; intros [r Hr]; exists r; simpl in *; assumption.
  Qed.

  #[global] Instance thunkGenUnsized {A} (f : unit -> G A)
          `{@Unsized _ _ ProducerGen (f tt)} : Unsized (thunkGen f).
  Proof.
    intros s1 s2.
    do 2 rewrite semThunkGenSize.
    apply unsized.
  Qed.

  #[global] Instance thunkGenSizeMonotonic {A} (f : unit -> G A)
          `{@SizeMonotonic _ _ ProducerGen (f tt)} : SizeMonotonic (thunkGen f).
  Proof.
    intros s1 s2 Hs.
    do 2 rewrite semThunkGenSize.
    by apply monotonic.
  Qed.

  #[global] Instance thunkGenSizeMonotonicOpt {A} (f : unit -> G (option A))
          `{@SizeMonotonicOpt _ _ ProducerGen (f tt)} : SizeMonotonicOpt (thunkGen f).
  Proof.
    intros s1 s2 Hs. unfold semProdSizeOpt.
    do 2 rewrite semThunkGenSize.
    by apply monotonicOpt.
  Qed.

  (* #[global] Instance thunkGenSizeAntiMonotonicNone {A} (f : unit -> G (option A)) *)
  (*         `{@SizedAntimonotonicNone _ _ ProducerGen (f tt)} : SizedAntimonotonicNone (thunkGen f). *)
  (* Proof. *)
  (*   intros s1 s2 Hs. *)
  (*   do 2 rewrite semThunkGenSize. *)
  (*   by apply monotonicNone. *)
  (* Qed. *)

Fixpoint pick {A : Type} (def : G A) (xs : list (nat * G A)) n : nat * G A :=
  match xs with
    | nil => (0, def)
    | (k, x) :: xs =>
      if (n < k) then (k, x)
      else pick def xs (n - k)
  end.

(* This should use urns! *)
Fixpoint pickDrop {A : Type} (xs : list (nat * G (option A))) n : nat * G (option A) * list (nat * G (option A)) :=
  match xs with
    | nil => (0, ret None, nil)
    | (k, x) :: xs =>
      if (n < k) then  (k, x, xs)
      else let '(k', x', xs') := pickDrop xs (n - k)
           in (k', x', (k,x)::xs')
  end. 

Definition sum_fst {A : Type} (gs : list (nat * A)) : nat :=
  foldl (fun t p => t + (fst p)) 0 gs.

Definition freq_ {A : Type} (def : G A) (gs : list (nat * G A))
: G A :=
  let tot := sum_fst gs in
  bindGen (choose (0, tot-1)) (fun n =>
  @snd _ _ (pick def gs n)).

(*
Definition frequency {A}:= 
  @deprecate (G A -> list (nat * G A) -> G A) "frequency" "freq_" freq_.
 *)

Fixpoint backtrackFuel {A : Type} (fuel : nat) (tot : nat) (gs : list (nat * G (option A))) : G (option A) :=
  match fuel with 
    | O => ret None
    | S fuel' => bindGen (choose (0, tot-1)) (fun n => 
                 let '(k, g, gs') := pickDrop gs n in
                 bindGen g (fun ma =>
                 match ma with 
                   | Some a => ret (Some a)
                   | None => backtrackFuel fuel' (tot - k) gs'
                 end ))
  end.

Definition backtrack {A : Type} (gs : list (nat * G (option A))) : G (option A) :=
  backtrackFuel (length gs) (sum_fst gs) gs.

Definition retryBody {A : Type}
           (retry : nat -> G (option A) -> G (option A))
           (n : nat) (g : G (option A)) : G (option A) :=
  bindGen g (fun x =>
               match x, n with
               | Some a, _ => returnGen (Some a)
               | None, O => returnGen None
               | None, S n' => retry n' g
               end).

(* Rerun a generator [g] until it returns a [Some], or stop after
     [n+1] tries. *)
Fixpoint retry {A : Type} (n : nat) (g : G (option A)) :
  G (option A) :=
  retryBody retry n g.

(* Filter the output of a generator [g], returning [None] when the
     predicate [p] is [false]. The generator is run once. *)
Definition suchThatMaybe1 {A : Type} (g : G A) (p : A -> bool) :
  G (option A) :=
  fmap (fun x => if p x then Some x else None) g.

(* Retry a generator [g : G A] until it returns a value satisfying
     the predicate, or stop after [size+1] times, where [size] is the
     current size value. *)
Definition suchThatMaybe {A : Type} (g : G A) (p : A -> bool) :
  G (option A) :=
  sized (fun n => retry n (suchThatMaybe1 g p)).

(* Retry a generator [g : G (option A)] until it returns a value
     satisfying the predicate, or stop after [size+1] times, where
     [size] is the current size value. *)
Definition suchThatMaybeOpt {A : Type} (g : G (option A))
           (p : A -> bool) : G (option A) :=
  sized (fun n => retry n (fmap (fun x =>
                                   match x with
                                   | None => None
                                   | Some a => if p a then Some a else None
                                   end) g)).

(* Retry a generator until it returns a value, or stop after
     [size+1] times. *)
Definition retrySized {A : Type} (g : G (option A)) : G (option A) :=
  sized (fun n => retry n g).

(* begin semReturn *)
Lemma semReturnGen {A} (x : A) : semProd (ret x) <--> [set x].
(* end semReturn *)
Proof.
  rewrite /semProd /semProdSize /= /semGenSize /=
  bigcup_const ?codom_const //.
  exact: randomSeed_inhabited.
    by do 2! constructor.
Qed.
  
Lemma semReturnSizeGen A (x : A) (s : nat) :
  semProdSize (ret x) s <--> [set x].
Proof.
  rewrite /semProdSize /= /semGenSize.
  rewrite codom_const; [ reflexivity | apply randomSeed_inhabited ].
Qed.

Lemma randomSplit_codom : codom randomSplit <--> setT.
Proof.
by apply/subset_eqP; split=> // [[s1 s2]] _; apply: randomSplitAssumption.
Qed.


Lemma semBindSizeGen A B (g : G A) (f : A -> G B) (s : nat) :
    semGenSize (bindGen g f) s <-->
    \bigcup_(a in semGenSize g s) semGenSize (f a) s.
Proof.
    rewrite /semGenSize /bindGen /= bigcup_codom -curry_codom2l.
    by rewrite -[codom (uncurry _)]imsetT -randomSplit_codom -codom_comp.
Qed.

Lemma semSizedGen A (f : nat -> G A) :
    semProd (sized f) <--> \bigcup_n semGenSize (f n) n.
Proof. by []. Qed.

Lemma semSizedSizeGen A (f:nat->G A)s : semGenSize (sized f) s <--> semGenSize (f s) s.
Proof. by []. Qed.

Lemma semResizeGen A n (g : G A) : semProd (resize n g) <--> semGenSize g n .
Proof.
   by case: g => g; rewrite /semProd /semProdSize /= /semGenSize /= bigcup_const.
Qed.

Lemma semResizeSizeGen A (s n : nat) (g : G A) :
    semGenSize (resize n g) s <--> semGenSize g n.
Proof. by case: g => g; rewrite /semGenSize. Qed.

#[global] Instance ProducerSemanticsGen :
  @ProducerSemantics G ProducerGen :=
  {
  semReturn     := @semReturnGen; 
  semReturnSize := @semReturnSizeGen;
  semBindSize   := @semBindSizeGen;
  semChoose     := @semChooseGen;
  semChooseSize := @semChooseSizeGen;
  semSized      := @semSizedGen;
  semSizedSize  := @semSizedSizeGen;
  semResize     := @semResizeGen;
  semResizeSize := @semResizeSizeGen
  }.

Module QcDefaultNotation.

Declare Scope qc_scope.

Notation " 'elems' [ x ] " := (elems_ x (cons x nil)) : qc_scope.
Notation " 'elems' [ x ; y ] " := (elems_ x (cons x (cons y nil))) : qc_scope.
Notation " 'elems' [ x ; y ; .. ; z ] " :=
  (elems_ x (cons x (cons y .. (cons z nil) ..))) : qc_scope.
Notation " 'elems' ( x ;; l ) " :=
  (elems_ x (cons x l)) (at level 1, no associativity) : qc_scope.

Notation " 'oneOf' [ x ] " := (oneOf_ x (cons x nil)) : qc_scope.
Notation " 'oneOf' [ x ; y ] " := (oneOf_ x (cons x (cons y nil))) : qc_scope.
Notation " 'oneOf' [ x ; y ; .. ; z ] " :=
  (oneOf_ x (cons x (cons y .. (cons z nil) ..))) : qc_scope.
Notation " 'oneOf' ( x ;; l ) " :=
  (oneOf_ x (cons x l))  (at level 1, no associativity) : qc_scope.

Notation " 'freq' [ x ] " := (freq_ x nil) : qc_scope.
Notation " 'freq' [ ( n , x ) ; y ] " :=
  (freq_ x (cons (n, x) (cons y nil))) : qc_scope.
Notation " 'freq' [ ( n , x ) ; y ; .. ; z ] " :=
  (freq_ x (cons (n, x) (cons y .. (cons z nil) ..))) : qc_scope.
Notation " 'freq' ( ( n , x ) ;; l ) " :=
  (freq_ x (cons (n, x) l)) (at level 1, no associativity) : qc_scope.

End QcDefaultNotation.

Section FrequencyProof.
(* A rather long frequency proof, probably we can do better *)

Lemma not_lt : forall n m, (false = (n < m)) -> (m <= n).
Proof.
  move => n m. by elim: n m=> [| n IHn]; case.
Qed.

Lemma sum_fstE A x (a : A) l:
  sum_fst ((x, a) :: l) = x + sum_fst l.
Proof.
rewrite /sum_fst /=.
elim: l 0 x => [n x|[n1 x1] l IHl p q] /=; first by rewrite addnC.
by rewrite -IHl; congr foldl; rewrite addnAC.
Qed.

Lemma sum_fst_eq0P {A} (l : list (nat * A)) :
  sum_fst l = 0 <->
  [seq x <- l | x.1 != 0] = [::].  
Proof.
by elim: l => [|[[|n] x] l IHl] //=; split=> //; rewrite sum_fstE.
Qed.

Lemma pick_def :
  forall {A} (l: list (nat * G A)) n def,
    sum_fst l <= n ->
    pick def l n = (0, def).
Proof.
  move=> A l n def Hleq.
  elim : l n Hleq => //=. case=> //= i p l IHl n Hleq.
  rewrite sum_fstE in Hleq.
  remember (n < i). case: b Heqb => Heqb; symmetry in Heqb.
  - have : (i + sum_fst l) < i by eapply (leq_ltn_trans); eassumption.
    rewrite -ltn_subRL. by have -> : forall i, (i - i) = 0 by elim.
  - apply IHl. rewrite -(leq_add2r i) subnK.
      by rewrite addnC. by apply/not_lt.
Qed.

Lemma pick_exists :
  forall {A} (l: list (nat * G A)) n def,
    n <  sum_fst l <->
    exists x, List.In x l /\ pick def l n = x /\ fst x <> 0.
Proof.
  move => A l n def. split.
  - move => Hlt.
    elim : l n Hlt => //. case => i p xs IHxs n Hlt.
    rewrite sum_fstE in Hlt.
    move/(_ (n-i)) : IHxs => IHxs. simpl.
    remember (n < i). case: b Heqb => [Heqb | /not_lt Heqb].
    + exists (i, p). split => //=. by left.  split => //=.
      move => contra; subst. by rewrite ltn0 in Heqb.
    + rewrite -(ltn_add2r i) [X in _  < X]addnC subnK // in IHxs.
      move/(_ Hlt) : IHxs => [x [H1 [H2 H3]]].
      by exists x; split; [right | split].
  - move => [x [HIn [Hpick Hneq]]].
    remember (n < sum_fst l).
    case: b  Heqb => //= /not_lt/pick_def H.
    rewrite H in Hpick. rewrite -Hpick //= in Hneq.
Qed.

Lemma pick_In :
  forall {A} (l: list (nat * G A)) x def,
    List.In x l /\ fst x <> 0 ->
    exists n, pick def l n = x.
Proof.
  move => A l x def [HIn Hfst].
  elim : l HIn => //=. case => //= i g xs IHxs [H1 | H2]; subst.
  + exists 0. simpl in *.
    have H : 0 < i by  elim : i Hfst IHxs => //=.
    rewrite H.
      by split => //=.
  + move/(_ H2) : IHxs => [n Hpick].
    exists (n + i). rewrite -[X in _ < X]add0n ltn_add2r ltn0.
      by rewrite  -[X in _ - X]add0n subnDr subn0.
Qed.

Lemma pick_imset A (def : G A) l :
  pick def l @: [set m | m < sum_fst l] <--> [seq x <- l | x.1 != 0]. 
Proof.
elim: l => [|[n x] l IHl] /=.
  rewrite /sum_fst /=.
  have->: (fun m => m < 0) <--> set0 by [].
  by rewrite imset0.
case: n => /= [|n].
  rewrite -IHl => t; split=> [[y []]|].
    by rewrite sum_fstE add0n subn0 => lt_y <-; exists y.
  by case=> y [lt_y <-]; exists y; split=> //; rewrite subn0.
move=> t; split=> /= [[p [lt_p]]|].
  case: ifP => [_ <-|lt_pn ?]; first by left.
    right; rewrite -(IHl t); exists (p - n.+1); split=> //.
  rewrite sum_fstE in lt_p.
  by rewrite -(ltn_add2r n.+1) subnK 1?addnC // leqNgt lt_pn.
case=> [<-|]; first by exists 0; split => //; rewrite sum_fstE.
rewrite -(IHl t); case=> p [lt_p <-]; exists (n.+1 + p); split.
  by rewrite sum_fstE ltn_add2l.
by rewrite ltnNge leq_addr addKn.
Qed.

Lemma pickDrop_def :
  forall {A} (l: list (nat * G (option A))) n,
    sum_fst l <= n ->
    pickDrop l n = (0, ret None, l).
Proof.
  move=> A l n Hleq.
  elim : l n Hleq => //=. case=> //= i p l IHl n Hleq.
  rewrite sum_fstE in Hleq.
  remember (n < i). case: b Heqb => Heqb; symmetry in Heqb.
  - have : (i + sum_fst l) < i by eapply (leq_ltn_trans); eassumption.
    rewrite -ltn_subRL. by have -> : forall i, (i - i) = 0 by elim.
  - rewrite IHl; auto. rewrite -(leq_add2r i) subnK.
      by rewrite addnC. by apply/not_lt.
Qed.

(* Probably needs something about l' and l. *)
Lemma pickDrop_exists :
  forall {A} (l: list (nat * G (option A))) n,
    n <  sum_fst l <->
    exists k g l',
      List.In (k,g) l /\ pickDrop l n = (k,g,l') /\ k <> 0 /\
      l <--> [set (k, g)] :|: l' /\
      length l' + 1 = length l /\
      sum_fst l' + k = sum_fst l.
Proof.
  move => A l n. split.
  - move => Hlt.
    elim : l n Hlt => //. case => i p xs IHxs n Hlt.
    rewrite sum_fstE in Hlt.
    move/(_ (n-i)) : IHxs => IHxs. simpl.
    remember (n < i). case: b Heqb => [Heqb | /not_lt Heqb].
    + exists i. exists p. exists xs. split => //=. by left.  split => //=.
      split. move => contra; subst. by rewrite ltn0 in Heqb.
      split. by rewrite cons_set_eq.
      split. by ssromega.
      rewrite sum_fstE. by ssromega.
    + rewrite -(ltn_add2r i) [X in _  < X]addnC subnK // in IHxs.
      move/(_ Hlt) : IHxs => [k [g [gs [H1 [H2 [H3 [H4 [H5 H6]]]]]]]].
      exists k. exists g. exists ((i,p)::gs).
      split; [right | split; [| split; [| split; [| split]]]];
      try (simpl; eauto; by rewrite H2).
      rewrite !cons_set_eq H4.
      rewrite setU_assoc (setU_comm [set (i, p)]) -setU_assoc.
      reflexivity.
      simpl. by ssromega.
      simpl. rewrite !sum_fstE. by ssromega.
  - move => [k [g [gs [HIn [Hpick [Hneq _]]]]]].
    remember (n < sum_fst l).
    case: b  Heqb => //= /not_lt/pickDrop_def H.
    rewrite H in Hpick. 
    inversion Hpick; subst; eauto.
Qed.

Lemma pickDrop_In :
  forall {A} (l: list (nat * G (option A))) k x,
    List.In (k,x) l /\ k <> 0 ->
    exists n l', pickDrop l n = (k,x,l').
Proof.
  move => A l k x [HIn Hfst].
  elim : l HIn => //=. case => //= i g xs IHxs [H1 | H2]; subst.
  + exists 0.  exists xs. simpl in *.
    inversion H1; subst; clear H1.
    have H : 0 < k by  elim : k Hfst IHxs => //=.
    rewrite H.
      by split => //=.
  + move/(_ H2) : IHxs => [n [l' Hpick]].
    exists (n + i). exists ((i,g)::l'). 
    rewrite -[X in _ < X]add0n ltn_add2r ltn0.
    rewrite  -[X in _ - X]add0n subnDr subn0.
    by rewrite Hpick.
Qed.

Lemma pickDrop_In_strong :
  forall {A} (l: list (nat * G (option A))) k x,
    List.In (k,x) l /\ k <> 0 ->
    exists n l',
      pickDrop l n = (k,x,l') /\
      n < sum_fst l /\
      length l = length l' + 1.
Proof.
  move => A l k x [HIn Hfst].
  elim : l HIn => //=. case => //= i g xs IHxs [H1 | H2]; subst.
  + exists 0.  exists xs. simpl in *.
    inversion H1; subst; clear H1.
    have H : 0 < k by  elim : k Hfst IHxs => //=.
    rewrite H. split ; [| split ]; simpl; auto.
    rewrite sum_fstE. now ssromega.
    now ssromega.
  + move/(_ H2) : IHxs => [n [l' [Hpick [Hlt Hlen]]]].
    exists (n + i). exists ((i,g)::l'). 
    rewrite -[X in _ < X]add0n ltn_add2r ltn0.
    rewrite  -[X in _ - X]add0n subnDr subn0.
    rewrite Hpick. simpl.
    split ; [| split ]; simpl; auto.
    rewrite sum_fstE. now ssromega.
    now ssromega.
Qed.

(* begin semFrequencySize *)
Lemma semFrequencySize {A}
      (l : list (nat * G A)) (def : G A) (size: nat) :
  semProdSize (freq_ def l) size <-->
    let l' := [seq x <- l | x.1 != 0] in
    if l' is nil then semProdSize def size else
      \bigcup_(x in l') semProdSize x.2 size.
(* end semFrequencySize *)
Proof.
rewrite semBindSize semChooseSizeGenNat //=.
case lsupp: {1}[seq x <- l | x.1 != 0] => [|[n g] gs].
move/sum_fst_eq0P: lsupp => suml; rewrite suml.
  rewrite (@eq_bigcupl _ _ _ [set 0]) ?bigcup_set1 ?pick_def // ?leqn0 ?suml //.
  by move=> n; split; rewrite leqn0; [move/eqP|] => ->.
symmetry; apply: reindex_bigcup.
have pos_suml: 0 < sum_fst l.
  have [] := sum_fst_eq0P l.
  by rewrite lsupp; case: (sum_fst l) => // /(_ erefl).
have->: (fun a : nat => a <= sum_fst l - 1) <--> [set m | m < sum_fst l].
  by move=> m /=; rewrite -ltnS subn1 prednK.
exact: pick_imset.
Qed.

(* begin semFrequency *)
Lemma semFrequency {A} (l : list (nat * G A)) (def : G A) :
  semProd (freq_ def l) <-->
    let l' := [seq x <- l | x.1 != 0] in
    if l' is nil then semProd def else
      \bigcup_(x in l') semProd x.2.
(* end semFrequency *)
Proof.
by case lsupp: {1}[seq x <- l | x.1 != 0] => [|[n g] gs] /=;
rewrite 1?bigcupC; apply: eq_bigcupr => sz;
have := (semFrequencySize l def sz); rewrite lsupp.
Qed.

Lemma frequencySizeMonotonic {A} (g0 : G A) lg :
  SizeMonotonic g0 ->
  List.Forall (fun p => SizeMonotonic (snd p)) lg ->
  SizeMonotonic (freq_ g0 lg).
Proof.
  intros H1.  unfold freq_.
  intros Hall. eapply bindMonotonicStrong.
  - eauto with typeclass_instances.
  - apply unsizedMonotonic.
    apply chooseUnsized.
  - intros x Heq. eapply semChooseGenNat in Heq; eauto.
  move : Heq => /andP [Hep1 Heq2]. 
  destruct (sum_fst lg) eqn:Heq.
  + rewrite pick_def. eassumption.
    subst. ssromega.
  + edestruct (pick_exists lg x g0) as [[[n' g] [Hin [Hp Hg]]] H2].
    rewrite Heq. ssromega.
    eapply List.Forall_forall in Hall; [ | ].
    eassumption.
    subst. rewrite Hp. eassumption.
Qed.

#[global] Instance frequencySizeMonotonic_alt :
  forall {A : Type} (g0 : G A) (lg : seq (nat * G A)),
    SizeMonotonic g0 ->
    lg \subset [set x | SizeMonotonic x.2 ] ->
    SizeMonotonic (freq_ g0 lg).
Proof.
  intros A g ls Hm Hin.
  eapply frequencySizeMonotonic. eassumption.
  induction ls. now constructor.
  constructor. eapply Hin.
  constructor. reflexivity.
  eapply IHls.  eapply subset_trans; eauto.
  constructor 2. eassumption.
Qed.

End FrequencyProof.  

Lemma backtrack_correct_size_opt {A} (lst : list (nat * G (option A))) s :
  semProdSizeOpt (backtrack lst) s <-->
  \bigcup_(x in lst :&: (fun x => x.1 <> 0)) (fun g => semProdSizeOpt (snd g) s) x.  
Proof.
  unfold backtrack.
  assert (Hret := @semReturnSize G _ _ (option A)).
  assert (Hbind := @semBindSize G _ _). simpl in *. 
  

  assert (sum_fst lst = sum_fst lst)%coq_nat by reflexivity.  
  revert H.
  assert (Datatypes.length lst = Datatypes.length lst)%coq_nat by reflexivity.  
  revert H.
  
  generalize (sum_fst lst) at 1 3.
  generalize (Datatypes.length lst) at 1 3.
  
  intros n1. generalize lst. induction n1; intros l n2 Heq1 Heq2.
  - simpl. intros x; split; intros Hin.
    + eapply semReturnSizeOpt_None in Hin; eauto with typeclass_instances. inv Hin.
    + inv Hin. inv H. destruct l; try (simpl in *; congruence). inv H0.
      inv H2.
  - intros x; split; intros Hin.
    + with_strategy opaque [pickDrop] (simpl in Hin). 
      eapply Hbind in Hin.
      inv Hin. inv H.
      eapply semChooseSizeGen in H0; eauto. simpl in *.
      
    (*   destruct (pickDrop_exists l x). simpl in *. destruct H4.       *)

    (*   now lia. *)
    (*   destruct H1. destruct H4. destruct H4. destruct H6. destruct H7. *)
      
    (*   rewrite H4 in H3. *)
      
    (*   eapply Hbind in H3. *)
    (*   destruct H2. destruct H3. destruct H2. *)
    (*   destruct x2. *)
      
      
    (*   -- eapply Hret in H3. *)
    (*      inv H3. *)
         
    (*      eexists. split.  eassumption. *)
    (*      constructor; eauto. *)
         
    (*   -- assert (Hsem : (isSome :&: semProdSize *)
    (*                             (enumerateFuel n (n.+1 - 1) *)
    (*                                            x1) s) (Some a)). *)
    (*      { split; eauto. } *)
         
    (*      assert (Heq' : (n.+1 - 1) = n). *)
    (*      { ssromega. } *)
         
    (*      rewrite Heq' in Hsem. *)
    (*      eapply IHn in Hsem. *)
    (*      inv Hsem. destruct H9. *)
    (*      eexists. split. *)
    (*      eapply H8. eassumption.  *)
    (*      eassumption. *)
    (*      ssromega.  *)
    (* + inv Hin. inv H. inv H1. destruct x; try (now exfalso; eauto). *)
    (*   constructor. now eauto. *)
    (*   simpl.  *)
    (*   eapply Hbind. *)
      
    (*   destruct (pickDrop_In _ _ H0). destruct H4. *)
    (*   destruct H4. *)
      
    (*   exists x. split. *)
    (*   eapply Enumerators.semChooseSize; eauto. *)
    (*   simpl. now ssromega.  *)

    (*   rewrite H4. *)
    (*   eapply Hbind. *)

    (*   exists (Some a). split. *)
    (*   eassumption. *)
    (*   eapply Hret. reflexivity.  *)
Admitted. (* TODO bring back to life *)       


Lemma backtrack_correct_opt {A} (lst : list (nat * G (option A))) :
  semProdOpt (backtrack lst) <--> \bigcup_(x in lst :&: fun x => x.1 <> 0) (semProdOpt x.2). 
Proof.
  split; intros H.
  - inv H. inv H0.
    assert (Hin : semProdSizeOpt (backtrack lst) x a).
    { eassumption. }
    eapply (@backtrack_correct_size_opt A) in Hin.
    inv Hin. inv H3.
    eexists. split; eauto. eexists. split; eauto.
  - destruct H. destruct H. destruct H0. destruct H0.
    destruct x. simpl in *. inv H. simpl in *.
    assert (Hin :  (\bigcup_(x in lst :&: fun x => x.1 <> 0) (semProdSizeOpt x.2 x0)) a).
    { eexists. split; eauto. }
    eapply (@backtrack_correct_size_opt A) in Hin.
    eexists. split; eauto. 
Qed.

Lemma backtrack_SizeMonotonicOpt (A : Type) (l : list (nat * G (option A))) :
  l \subset (fun x => SizeMonotonicOpt x.2) ->
  SizeMonotonicOpt (backtrack l).
Proof.
  intros Hin. intros s1 s2 Hleq.
  rewrite !backtrack_correct_size_opt.
  intros x Hin'.  destruct Hin' as [e [Hl Hs]].
  eexists. split; eauto. eapply Hin; inv Hl; eauto.  
Qed.

Lemma enumerate_SizeMonotonic (A : Type) (l : list (nat * G (option A))) :
  l \subset (fun x => SizeMonotonic x.2) ->
  SizeMonotonic (backtrack l).
Proof.
(*   unfold backtrack.  *)
(*   assert (Datatypes.length l = Datatypes.length l)%coq_nat by reflexivity.   *)
(*   revert H. *)
(*   generalize (Datatypes.length l) at 2 3 4. *)
(*   intros n. revert l. induction n; intros l Heq Hsub. *)
(*   - simpl. now eauto with typeclass_instances. *)
(*   - simpl. *)
(*     eapply bindMonotonicStrong; eauto with typeclass_instances. *)
    
(*     intros x1 Hin. eapply Enumerators.semChoose in Hin; eauto. simpl in *.  *)
    
(*     destruct (Enumerators.pickDrop_exists l x1). simpl in *. now ssromega. *)
(*     destruct H. destruct H. destruct H0. destruct H1. *)

(*     rewrite H. *)

(*     eapply bindMonotonicStrong; eauto with typeclass_instances. *)
    
(*     intros a Hin'. *)

(*     destruct a; eauto with typeclass_instances.  *)

(*     eapply returnGenSizeMonotonic; eauto with typeclass_instances. *)

(*     assert (Heq' : (n.+1 - 1) = n). { ssromega. } *)

(*     rewrite Heq'. eapply IHn.  *)

(*     now ssromega. *)

(*     eapply subset_trans. eassumption. eassumption. *)
(* Qed. *)
Admitted. (* TODO bring back to life *)

(* Backwards compatibility. *)
Definition elements := @elems_ G ProducerGen.
Definition liftGen := @liftM G (@super _ ProducerGen).
Definition liftGen2 := @liftM2 G (@super _ ProducerGen).
Definition liftGen3 := @liftM3 G (@super _ ProducerGen).
Definition liftGen4 := @liftProd4 _ ProducerGen.
Definition liftGen5 := @liftProd5 _ ProducerGen.
Definition sequenceGen := @sequenceProd G ProducerGen.
Definition oneof := @oneOf_ G ProducerGen.
Definition frequency := @freq_.
Definition semGen := @semProd G ProducerGen.
