(* The definitions and other results of age_by and age_to should be
moved here from msl/ageable.v.  Alternatively, this can be moved to
msl/ageable.v (or this file to msl/) eventually, but we keep it here
for now to reduce compilation time. *)

Require Import Coq.omega.Omega.
Require Import msl.ageable.
Require Import msl.predicates_hered.
Require Import msl.sepalg.
Require Import msl.age_sepalg.
Require Import msl.sepalg_generators.

(* Apply [age1] n times (meaningful when [n <= level x] *)

Definition age1' {A} `{agA : ageable A} : A -> A :=
  fun x => match age1 x with Some y => y | None => x end.

Definition age_by n {A} `{agA : ageable A} : A -> A := Nat.iter n age1'.

Lemma level_age1' {A} `{agA : ageable A} x : level (age1' x) = level x - 1.
Proof.
  unfold age1'. destruct (age1 x) eqn:E.
  - apply age_level in E. omega.
  - apply age1_level0 in E. omega.
Qed.

Lemma level_age_by n {A} `{agA : ageable A} x : level (age_by n x) = level x - n.
Proof.
  revert x; induction n; intros x; simpl.
  - omega.
  - simpl. rewrite level_age1'. rewrite IHn. omega.
Qed.

Lemma age_age_by n {A} `{agA : ageable A} (x y : A) : age x y -> age_by (S n) x = age_by n y.
Proof.
  intros E.
  induction n.
  - simpl.
    unfold age1' in *.
    rewrite E. auto.
  - change (age1' (age_by (S n) x) = age_by (S n) y).
    rewrite IHn.
    reflexivity.
Qed.

(* Age [x] to the level [k] (meaningul when [k <= level x] *)
Definition age_to k {A} `{agA : ageable A} (x : A) : A := age_by (level x - k) x.

Lemma level_age_to k {A} `{agA : ageable A} x : k <= level x -> level (age_to k x) = k.
Proof.
  intros L. unfold age_to.
  rewrite level_age_by; omega.
Qed.

(* Proof techniques for age_to *)
Lemma age_to_lt k {A} `{agA : ageable A} (x : A) : k < level x -> exists y, age x y /\ age_to k x = age_to k y.
Proof.
  intros L.
  destruct (age1 x) as [y|] eqn:Ex; swap 1 2.
  - rewrite age1_level0 in Ex. omega.
  - exists y; split; auto.
    unfold age_to.
    pose proof age_level _ _ Ex as E.
    replace (level x - k) with (S (level y - k)) by omega.
    generalize (level y - k).
    clear E L.
    intros.
    apply age_age_by, Ex.
Qed.

Lemma age_to_ge k {A} `{agA : ageable A} (x : A) : k >= level x -> age_to k x = x.
Proof.
  intros E. unfold age_to.
  replace (level x - k) with 0 by auto with *.
  reflexivity.
Qed.

Lemma age_to_eq k {A} `{agA : ageable A} (x : A) : k = level x -> age_to k x = x.
Proof.
  intros ->; apply age_to_ge, le_refl.
Qed.

Lemma age_age_to n {A} `{agA : ageable A} x y : level x = S n -> age x y -> age_to n x = y.
Proof.
  intros E Y.
  assert (L : (n < level x)%nat) by omega.
  unfold age_to. rewrite E. replace (S n - n) with 1 by omega.
  simpl. unfold age1'. rewrite Y. reflexivity.
Qed.

Lemma age_by_ind {A} `{agA : ageable A} (P : A -> Prop) :
  (forall x y, age x y -> P x -> P y) ->
  forall x n, P x -> P (age_by n x).
Proof.
  intros IH x n.
  unfold age_by.
  induction n; intros Px.
  - auto.
  - simpl. unfold age1' at 1.
    destruct (age1 (Nat.iter n age1' x)) as [y|] eqn:Ey; auto.
    eapply IH; eauto.
Qed.

Lemma age_to_ind {A} `{agA : ageable A} (P : A -> Prop) :
  (forall x y, age x y -> P x -> P y) ->
  forall x n, P x -> P (age_to n x).
Proof.
  intros IH x n.
  apply age_by_ind, IH.
Qed.

Lemma age_to_ind_refined n {A} `{agA : ageable A} (P : A -> Prop) :
  (forall x y, age x y -> n <= level y -> P x -> P y) ->
  forall x, P x -> P (age_to n x).
Proof.
  intros IH x Px.
  assert (dec : n >= level x \/ n <= level x) by omega.
  destruct dec as [Ge|Le].
  - rewrite age_to_ge; auto.
  - eapply (age_to_ind (fun x => n <= level x -> P x)); auto.
    + intros x0 y H H0 H1.
      eapply IH; eauto.
      apply age_level in H.
      apply H0.
      omega.
    + rewrite level_age_to; auto.
Qed.

Lemma iter_iter n m {A} f (x : A) : Nat.iter n f (Nat.iter m f x) = Nat.iter (n + m) f x.
Proof.
  induction n; auto; simpl. rewrite IHn; auto.
Qed.

Lemma age_by_age_by n m  {A} `{agA : ageable A} (x : A) : age_by n (age_by m x) = age_by (n + m) x.
Proof.
  apply iter_iter.
Qed.

Lemma age_by_ind_opp {A} `{agA : ageable A} (P : A -> Prop) :
  (forall x y, age x y -> P y -> P x) ->
  forall x n, P (age_by n x) -> P x.
Proof.
  intros IH x n.
  unfold age_by.
  induction n; intros Px.
  - auto.
  - simpl in Px. unfold age1' at 1 in Px.
    destruct (age1 (Nat.iter n age1' x)) as [y|] eqn:Ey; auto.
    eapply IH in Ey; eauto.
Qed.

Lemma age_to_ind_opp {A} `{agA : ageable A} (P : A -> Prop) :
  (forall x y, age x y -> P y -> P x) ->
  forall x n, P (age_to n x) -> P x.
Proof.
  intros IH x n.
  apply age_by_ind_opp, IH.
Qed.

Lemma rewrite_age_to {A} `{agA : ageable A} (P : A -> Prop) :
  (forall x y, age x y -> P x <-> P y) ->
  forall x n, P x <-> P (age_to n x).
Proof.
  intros IH x n; split.
  - apply age_to_ind. intros; rewrite <-IH; eauto.
  - apply age_to_ind_opp. intros; rewrite IH; eauto.
Qed.

Lemma level_age_to_le k {A} `{agA : ageable A} x : level (age_to k x) <= level x.
Proof.
  destruct (le_lt_dec k (level x)) as [l|l]. rewrite level_age_to; auto.
  rewrite age_to_ge; omega.
Qed.

Lemma level_age_to_le' k {A} `{agA : ageable A} x : level (age_to k x) <= k.
Proof.
  destruct (le_lt_dec k (level x)) as [l|l]. rewrite level_age_to; auto.
  rewrite age_to_ge; omega.
Qed.

Lemma age_by_necR {A} `{agA : ageable A} n x : necR x (age_by n x).
Proof.
  generalize (necR_refl x).
  generalize x at 1 3; intros u.
  apply age_by_ind; clear x.
  intros x y a N.
  constructor 3 with x; auto.
  constructor; auto.
Qed.

Lemma age_to_necR {A} `{agA : ageable A} n x : necR x (age_to n x).
Proof.
  apply age_by_necR.
Qed.

Lemma necR_age_by {A} `{agA : ageable A} x x' : necR x x' -> x' = age_by (level x - level x') x.
Proof.
  intros N; induction N.
  - rewrite (age_level _ _ H).
    replace (S _ - _) with 1. 2:omega.
    simpl. unfold age1'. rewrite H; auto.
  - replace (_ - _) with 0. 2:omega. reflexivity.
  - rewrite IHN2, IHN1.
    rewrite age_by_age_by.
    repeat rewrite level_age_by.
    f_equal.
    apply necR_level in N1.
    apply necR_level in N2.
    replace (_ x - (_ x - _ y)) with (level y) by omega.
    replace (_ y - _ z + (_ x - _ y)) with (level x - level z) by omega.
    omega.
Qed.

Lemma necR_age_to {A} `{agA : ageable A} x x' : necR x x' -> x' = age_to (level x') x.
Proof.
  apply necR_age_by.
Qed.

Lemma necR_age_by_iff {A} `{agA : ageable A} x x' : necR x x' <-> x' = age_by (level x - level x') x.
Proof.
  split. apply necR_age_by. intros ->. apply age_by_necR.
Qed.

Lemma necR_age_to_iff {A} `{agA : ageable A} x x' : necR x x' <-> x' = age_to (level x') x.
Proof.
  apply necR_age_by_iff.
Qed.

Lemma age_to_pred {A} `{agA : ageable A} (P : pred A) x n :
  app_pred P x ->
  app_pred P (age_to n x).
Proof.
  apply age_to_ind. clear x n.
  destruct P as [x h]. apply h.
Qed.

Lemma age_by_pred {A} `{agA : ageable A} (P : pred A) x n :
  app_pred P x ->
  app_pred P (age_by n x).
Proof.
  apply age_by_ind. clear x n.
  destruct P as [x h]. apply h.
Qed.

Lemma pred_age1' {A} `{agA : ageable A} (R : pred A) x : app_pred R x -> app_pred R (age1' x).
Proof.
  unfold age1'. destruct (age1 x) as [phi' | ] eqn:Ephi'; auto.
  destruct R as [R h]. apply h. apply Ephi'.
Qed.

Lemma age_by_age_by_pred {A} `{agA : ageable A} (P : pred A) x n1 n2 :
  le n1 n2 ->
  app_pred P (age_by n1 x) ->
  app_pred P (age_by n2 x).
Proof.
  intros l. replace n2 with ((n2 - n1) + n1) by auto with *.
  rewrite <-age_by_age_by.
  apply age_by_pred.
Qed.

Fixpoint composeOptN' {A} (f : A -> option A) n x :=
  match n with
  | 0 => Some x
  | S n =>
    match composeOptN' f n x with
    | Some y => f y
    | None => None
    end
  end.

Lemma composeOptN_assoc_aux_None {A} (f : A -> option A) n x :
  f x = None ->
  match composeOptN f n x with
  | Some x => f x
  | None => None
  end = None.
Proof.
  revert x; induction n; intros x E; simpl; auto.
  destruct (f x); congruence.
Qed.

Lemma composeOptN_assoc_aux_Some {A} (f : A -> option A) n x y :
  f x = Some y ->
  match composeOptN f n x with
  | Some x => f x
  | None => None
  end = composeOptN f n y.
Proof.
  revert x y; induction n; intros x y Ey; simpl. auto.
  rewrite Ey.
  destruct (f y) as [z|] eqn:Ez.
  - eauto.
  - apply composeOptN_assoc_aux_None, Ez.
Qed.

Lemma composeOptN_assoc {A} (f : A -> option A) n x :
  composeOptN f n x = composeOptN' f n x.
Proof.
  revert x; induction n; intros x; simpl. auto.
  destruct (f x) as [y|] eqn:Ey; rewrite <-IHn.
  - erewrite composeOptN_assoc_aux_Some; eauto.
  - rewrite composeOptN_assoc_aux_None; eauto.
Qed.

Lemma age_by_ageN {A} `{agA : ageable A} n (x : A) :
  n <= level x ->
  ageN n x = Some (age_by n x).
Proof.
  revert x; induction n; intros x l. reflexivity.
  unfold ageN.
  rewrite composeOptN_assoc; simpl; rewrite <-composeOptN_assoc.
  change (composeOptN age1 n x) with (ageN n x).
  rewrite IHn. 2:omega.
  unfold age1' in *.
  destruct (age1 (age_by n x)) as [y|] eqn:Ey. auto.
  exfalso. rewrite age1_level0 in Ey.
  rewrite level_age_by in Ey. omega.
Qed.

Lemma age_to_ageN {A} `{agA : ageable A} n (x : A) :
  ageN (level x - n) x = Some (age_to n x).
Proof.
  apply age_by_ageN. auto with *.
Qed.

Lemma age_by_1 {A} {_ : ageable A} x : 0 < level x -> age x (age_by 1 x).
Proof.
  intros l.
  unfold age_by, age1'; simpl.
  destruct (age1 x) eqn:E; eauto.
  apply age1_level0 in E.
  omega.
Qed.

Lemma age_to_1 {A} {_ : ageable A} n x : level x = S n -> age x (age_to n x).
Proof.
  unfold age_to; intros E; rewrite E.
  replace (S n - n) with 1 by auto with *.
  apply age_by_1. auto with *.
Qed.

Lemma age_to_identy {A} `{asaA: Age_alg A}: forall k phi,
    identity phi -> identity (age_to k phi).
Proof.
  intros k phi. unfold age_to. generalize (level phi - k); intros n; revert phi.
  induction n; intros phi id; simpl; auto. unfold age1'.
  destruct (age1 (age_by n phi)) eqn:E; auto.
  eapply age_identity. apply E. auto.
Qed.

Lemma age_to_join_eq {A}  {JA: Join A}{PA: Perm_alg A}{agA: ageable A}{AgeA: Age_alg A} :
  forall k x1 x2 x3,
    join x1 x2 x3 ->
    k <= level x3 ->
    join (age_to k x1) (age_to k x2) (age_to k x3).
Proof.
  intros k x1 x2 x3 J.
  remember (level x3) as l3 eqn:e3; symmetry in e3.
  pose proof join_level _ _ _ J as [e1 e2]; rewrite e3 in e1, e2.
  revert l3 x1 x2 x3 e1 e2 e3 J.
  intros n; induction n as [ n IHn ] using strong_nat_ind. intros x1 x2 x3 e1 e2 e3 J L.
  destruct (le_lt_eq_dec _ _ L) as [Lt | ->]; swap 1 2.
  now do 3 (rewrite age_to_eq at 1; auto).
  assert (l1 : k < level x1) by auto with *.
  assert (l2 : k < level x2) by auto with *.
  assert (l3 : k < level x3) by auto with *.
  destruct (age_to_lt _ x1 l1) as [x1' [E1 ->]].
  destruct (age_to_lt _ x2 l2) as [x2' [E2 ->]].
  destruct (age_to_lt _ x3 l3) as [x3' [E3 ->]].
  pose proof @age1_join_eq A _ _ _ _ _ _ _ J _ E1 _ E2 _ E3.
  pose proof @af_level2 A level age1 (@age_facts _ agA) _ _ E1.
  pose proof @af_level2 A level age1 (@age_facts _ agA) _ _ E2.
  pose proof @af_level2 A level age1 (@age_facts _ agA) _ _ E3.
  apply IHn with (level x1'); omega || auto.
Qed.


(* infotheo: information theory and error-correcting codes in Coq             *)
(* Copyright (C) 2020 infotheo authors, license: LGPL-2.1-or-later            *)
From mathcomp Require Import all_ssreflect fingroup perm.
From mathcomp Require boolp.
Require Import Reals.
From mathcomp Require Import Rstruct.
Require Import ssrR Reals_ext ssr_ext ssralg_ext logb Rbigop fdist entropy.
Require Import num_occ channel types.

(******************************************************************************)
(*                            Joint Types                                     *)
(*                                                                            *)
(* Definitions:                                                               *)
(*   jtype == type of joint type, notation: P_n(A, B)                         *)
(*   shell = shell, notation: V.-shell t                                      *)
(*   cond_type == conditional type, notation: \nu{V}(P)                       *)
(*                                                                            *)
(* Lemmas:                                                                    *)
(*   bound_card_jtype == upper-bound of the number of conditional types       *)
(*    jtype_not_empty == joint types are not empty                            *)
(*         occ_co_occ == Relation between the number of symbol occurrences    *)
(*                       and the number of pairs of symbols occurrences       *)
(******************************************************************************)

Reserved Notation "'P_' n '(' A ',' B ')'" (at level 9,
  n at next level, A at next level, B at next level).
Reserved Notation "V '.-shell' ta" (at level 5,
  format "V '.-shell'  ta").
Reserved Notation "'\nu_' n '^{' A ',' B '}' '(' P ')'" (at level 50,
  n, A, B, P at next level, format "'\nu_' n '^{' A ',' B '}' '(' P ')'").
Reserved Notation "'\nu^{' B '}' '(' P ')'" (at level 50,
  B, P at next level, format "'\nu^{' B '}' '(' P ')'").
Reserved Notation "'`tO(' V )" (at level 10).

Set Implicit Arguments.
Unset Strict Implicit.
Import Prenex Implicits.

Local Open Scope R_scope.
Local Open Scope channel_scope.

Module JType.
Section def.
Variables (A B : finType) (n : nat).
Local Open Scope nat_scope.

Record t : predArgType := mk {
  c :> `Ch*(A, B) ;
  f : {ffun A -> {ffun B -> 'I_n.+1}} ;
  sum_f : \sum_(a in A) \sum_(b in B) f a b == n ;
  c_f : forall a b, c a b = let row := \sum_(b in B) f a b in
                            if row == O
                            then / #|B|%:R
                            else ((f a b)%:R / row%:R)%R }.

End def.
End JType.

Coercion jtype_coercion := JType.c.

Notation "'P_' n '(' A ',' B ')'" := (JType.t A B n) : types_scope.
Local Open Scope types_scope.

Definition ffun_of_jtype A B n (i : P_ n ( A , B )) := let: JType.mk _ f _ _ := i in f.

Lemma jtype_ext A B n (t1 t2 : P_ n ( A , B )) : JType.f t1 = JType.f t2 -> t1 = t2.
Proof.
case: t1 t2 => d1 f1 Hf1 H1 /= [] d2 f2 Hf2 H2 /= ?; subst f2.
have ? : d1 = d2.
  apply/Channel1.chan_star_eq.
  by rewrite boolp.funeqE => a; apply/fdist_ext => b; rewrite H1 H2.
subst d2; congr JType.mk; exact/boolp.Prop_irrelevance.
Qed.

Definition jtype_eq A B n (t1 t2 : P_ n ( A , B )) :=
  match t1, t2 with
    | JType.mk _ f1 _ _, JType.mk _ f2 _ _ => f1 == f2
  end.

Lemma jtype_eqP A B n : Equality.axiom (@jtype_eq A B n).
Proof.
case=> d1 f1 Hf1 H1 [] d2 f2 Hf2 H2 /=.
apply: (iffP idP) => [/eqP H |[] _ -> //].
subst f2.
have ? : d1 = d2.
  apply/Channel1.chan_star_eq.
  by rewrite boolp.funeqE => a /=; apply/fdist_ext => b; rewrite H1 H2.
subst d2; congr JType.mk; exact/boolp.Prop_irrelevance.
Qed.

Definition jtype_eqMixin A B n := EqMixin (@jtype_eqP A B n).
Canonical jtype_eqType A B n := Eval hnf in EqType _ (@jtype_eqMixin A B n).

Definition nneg_fun_of_pre_jtype (A B : finType) (Bnot0 : (0 < #|B|)%nat) n
  (f : {ffun A -> {ffun B -> 'I_n.+1}}) : A -> nneg_finfun B.
pose pf := fun a => [ffun b : B =>
  let ln := (\sum_(b1 in B) (f a b1))%nat in
  if ln == O
    then / #|B|%:R
    else (f a b)%:R / ln%:R].
move=> a.
refine (@mkNNFinfun _ (pf a) _); apply/forallP_leRP => b.
rewrite /pf ffunE.
case: ifP => [_ | Hcase].
- exact/invR_ge0/ltR0n.
- apply divR_ge0; first exact/leR0n.
  by apply/ltRP; rewrite lt0R INR_eq0' Hcase /= leR0n'.
Defined.

Definition chan_of_jtype (A B : finType) (Anot0 : (0 < #|A|)%nat) (Bnot0 : (0 < #|B|)%nat)
  n (f : {ffun A -> {ffun B -> 'I_n.+1}}) : `Ch*(A, B).
set pf := fun a b =>
  let ln := (\sum_(b1 in B) (f a b1))%nat in
  if ln == O
  then / #|B|%:R
  else (f a b)%:R / ln%:R.
refine (@Channel1.mkChan A B _ Anot0) => a.
apply: (@FDist.mk _ (@nneg_fun_of_pre_jtype _ _ Bnot0 n f a)).
under eq_bigr do rewrite ffunE /=.
case/boolP : (\sum_(b1 in B) (f a b1) == O)%nat => Hcase.
- by rewrite /Rle big_const iter_addR mulRV // INR_eq0' -lt0n.
- rewrite big_morph_natRD /Rdiv -big_distrl /= mulRV //.
  by rewrite -big_morph_natRD // INR_eq0'.
Defined.

Definition jtype_choice_f (A B : finType) n (f : {ffun A -> {ffun B -> 'I_n.+1}}) : option (P_ n ( A , B )).
refine (
match Sumbool.sumbool_of_bool (0 < #|A|)%nat with
  | left Anot0 =>
    match Sumbool.sumbool_of_bool (0 < #|B|)%nat with
      |left Bnot0 =>
       match Sumbool.sumbool_of_bool (\sum_(a in A) \sum_(b in B) f a b == n)%nat with
         | left Hf => Some (@JType.mk A B n (chan_of_jtype Anot0 Bnot0 f) f Hf _)
         | right _ => None
       end
      | right _ => None
    end
  | right _ => None
end).
move=> a b; by rewrite ffunE.
Defined.

Lemma jtype_choice_pcancel (A B : finType) n : pcancel (@JType.f A B n) (@jtype_choice_f A B n).
Proof.
case=> d f Hf H /=.
rewrite /jtype_choice_f /=.
destruct d as [d Hd].
destruct Sumbool.sumbool_of_bool; last by rewrite Hd in e.
destruct Sumbool.sumbool_of_bool; last first.
  exfalso.
  case/card_gt0P : e => a Ha.
  move: e0.
  by rewrite (fdist_card_neq0 (d a)).
set d1 := chan_of_jtype _ _ _.
set d2 := Channel1.mkChan d Hd.
have d12 : d1 = d2.
  apply/Channel1.chan_star_eq.
  rewrite boolp.funeqE => /= a; apply/fdist_ext => b; by rewrite ffunE H.
destruct Sumbool.sumbool_of_bool; last by rewrite Hf in e1.
congr Some; by apply/jtype_eqP => /=.
Qed.

Lemma jtype_choiceMixin A B n : choiceMixin (P_ n ( A , B )).
Proof. apply (PcanChoiceMixin (@jtype_choice_pcancel A B n)). Qed.

Canonical jtype_choiceType A B n := Eval hnf in ChoiceType _ (jtype_choiceMixin A B n).

Definition jtype_pickle (A B : finType) n (P : P_ n (A , B)) : nat.
destruct P as [d t H].
exact: (pickle t).
(*destruct (finfun_countMixin A [finType of {ffun B -> 'I_n.+1}]) as [pi unpi Hcan].
apply pi.
exact t.*)
Defined.

Definition jtype_unpickle (A B : finType) n (m : nat) : option (P_ n ( A , B )).
refine (
    match Sumbool.sumbool_of_bool (0 < #|A|)%nat with
      | left Anot0 =>
        match Sumbool.sumbool_of_bool (0 < #|B|)%nat with
          |left Bnot0 => _
          | right _ => None
        end
      | right _ => None
    end
).
(*case: (finfun_countMixin A [finType of {ffun B -> 'I_n.+1}]) => pi unpi Hcan.
case: (unpi m); last by exact None.
move=> f.*)
pose unpi : option {ffun A -> {ffun B -> 'I_n.+1}} := unpickle m.
case: unpi; last first.
  exact None.
move=> f.
refine (match Sumbool.sumbool_of_bool (\sum_(a in A) \sum_(b in B) f a b == n)%nat with
          | left Hf => _
          | right _ => None
        end).
refine (Some (@JType.mk A B n (chan_of_jtype Anot0 Bnot0 f) f Hf _)).
by move=> a b; rewrite ffunE.
Defined.

Lemma jtype_count_pcancel A B n : pcancel (@jtype_pickle A B n) (@jtype_unpickle A B n).
Proof.
move=> V.
rewrite /jtype_unpickle /jtype_pickle /=.
destruct V as [[c Anot0] f Hf H].
destruct Sumbool.sumbool_of_bool; last by rewrite Anot0 in e.
case/card_gt0P : (e) => a Ha.
move: (fdist_card_neq0 (c a)) => Bnot0.
destruct Sumbool.sumbool_of_bool; last by rewrite Bnot0 in e0.
rewrite pickleK.
(*rewrite pcan_pickleK; last by apply valK.*)
set d1 := chan_of_jtype _ _ _.
have ? : d1 = Channel1.mkChan c Anot0.
  apply/Channel1.chan_star_eq.
  by rewrite boolp.funeqE => a1; apply/fdist_ext => b /=; rewrite ffunE H.
destruct Sumbool.sumbool_of_bool; last by rewrite Hf in e1.
congr Some; by apply/jtype_eqP => /=.
Qed.

Definition jtype_countMixin A B n := CountMixin (@jtype_count_pcancel A B n).

Canonical jtype_countType (A B : finType) n :=
  Eval hnf in CountType (P_ n ( A , B )) (@jtype_countMixin A B n).

Definition jtype_enum_f (A B : finType) n
  (f : { f : {ffun A -> {ffun B -> 'I_n.+1}} | (\sum_(a in A) \sum_(b in B) f a b == n)%nat}) :
  option (P_ n ( A , B )).
refine (
    match Sumbool.sumbool_of_bool (0 < #|A|)%nat with
      | left Anot0 =>
        match Sumbool.sumbool_of_bool (0 < #|B|)%nat with
          |left Bnot0 => _
          | right _ => None
        end
      | right _ => None
    end).
refine (Some (@JType.mk A B n (@chan_of_jtype _ _ Anot0 Bnot0 n (sval f)) (sval f) (proj2_sig f) _)).
by move=> a b; rewrite ffunE.
Defined.

Definition jtype_enum A B n := pmap (@jtype_enum_f A B n)
  (enum [finType of { f : {ffun A -> {ffun B -> 'I_n.+1}} | (\sum_(a in A) \sum_(b in B) f a b == n)%nat}]).

Lemma jtype_enumP A B n : Finite.axiom (@jtype_enum A B n).
Proof.
case=> d f Hf H /=.
have : Finite.axiom (enum [finType of { f : {ffun A -> {ffun B -> 'I_n.+1}}  |
    (\sum_(a in A) \sum_(b in B) f a b == n)%nat }]).
  rewrite enumT; by apply enumP.
move/(_ (@exist _ _ f Hf)) => <-.
rewrite /jtype_enum /=.
rewrite /jtype_enum_f /=.
destruct d as [d Anot0].
case/card_gt0P : (Anot0) => a _.
move: (fdist_card_neq0 (d a)) => Bnot0.
set tmp := pmap _ _.
have -> : tmp = pmap (fun f =>
                        Some {|
                            JType.c := chan_of_jtype Anot0 Bnot0 (sval f);
                            JType.f := (sval f);
                            JType.sum_f := proj2_sig f;
                            JType.c_f := fun a b =>
eq_ind_r
                                     (eq^~ (if (\sum_(b0 in B) sval f a b0)%nat == 0%N
                                            then / #|B|%:R
                                            else (sval f a b)%:R / (\sum_(b0 in B) sval f a b0)%:R))
                                     (erefl
                                        (if (\sum_(b0 in B) sval f a b0)%nat == 0%N
                                         then / #|B|%:R
                                         else (sval f a b)%:R / (\sum_(b0 in B) sval f a b0)%:R))
                                     (ffunE
                                        (fun b0 : B =>
                                         if (\sum_(b1 in B) sval f a b1)%nat == 0%N
                                         then / #|B|%:R
                                         else (sval f a b0)%:R / (\sum_(b1 in B) sval f a b1)%:R) b)
                                             (*(if \sum_(b0 in B) (sval f a) b0 == 0%N
                                              then / #|B|%:R
                                              else
                                                (sval f a b)%:R /
                                                    (\sum_(b0 in B) (sval f a) b0)%:R)*) |}
                     ) (enum [finType of { f : {ffun A -> {ffun B -> 'I_n.+1}} |
                                           (\sum_(a in A) \sum_(b in B) f a b)%nat == n}]).
  apply: eq_pmap => V.
  destruct Sumbool.sumbool_of_bool; last by rewrite Anot0 in e.
  destruct Sumbool.sumbool_of_bool; last by rewrite Bnot0 in e0.
  congr Some.
  by apply/jtype_eqP => /=.
rewrite count_map.
by apply eq_count.
Qed.

Definition jtype_finMixin A B n := Eval hnf in FinMixin (@jtype_enumP A B n).
Canonical jtype_finType A B n := Eval hnf in FinType _ (@jtype_finMixin A B n).

Section jtype_facts.
Variables (A B : finType) (n : nat) (ta : n.-tuple A).

Local Open Scope nat_scope.

Lemma jtype_entry_ub (V : P_ n ( A , B )) b : \sum_(a : A) (JType.f V) a b < n.+1.
Proof.
apply (@leq_ltn_trans (\sum_ (a : A) \sum_ (b : B) (JType.f V) a b)).
- apply leq_sum => a _.
  by rewrite (bigD1 b) //= leq_addr.
- case: V => c f Hf H /=.
  move/eqP in Hf.
  by rewrite Hf.
Qed.

Lemma jtype_0_jtypef (V : P_ n ( A , B )) a b : V a b = 0%R -> (JType.f V) a b = ord0.
Proof.
destruct V as [V1 V2 V3 V4] => /=.
rewrite V4 /=.
case: ifP => [| H'].
  rewrite sum_nat_eq0.
  move/forallP/(_ b)/implyP/(_ Logic.eq_refl)/eqP => H _; exact: val_inj.
rewrite /Rdiv mulR_eq0 => -[|abs].
  rewrite INR_eq0 => ?; exact/val_inj.
exfalso.
by apply/eqP : abs; apply/invR_neq0'; rewrite INR_eq0' H'.
Qed.

Lemma bound_card_jtype : #| P_ n (A , B) | <= expn n.+1 (#|A| * #|B|).
Proof.
rewrite -(card_ord n.+1) mulnC expnM -2!card_ffun cardE /enum_mem.
apply (@leq_trans (size (map (@ffun_of_jtype A B n) (Finite.enum (jtype_finType A B n))))).
  by rewrite 2!size_map.
rewrite cardE.
apply: uniq_leq_size.
  rewrite map_inj_uniq.
  rewrite -enumT; by apply enum_uniq.
  move=> [c f Hf Hc] [c1 f1 Hf1 Hc1] /= ?; subst f1.
  by apply/jtype_eqP => /=.
move=> /= x.
case/mapP => t Ht Hx.
subst x.
destruct t as [c f Hf cf] => /=.
apply/mapP.
by exists f.
Qed.

Lemma jtype_not_empty : 0 < #|A| -> 0 < #|B| -> 0 < #|P_ n (A , B)|.
Proof.
move=> Anot0 Bnot0.
move: (Anot0); case/card_gt0P => a _.
move: (Bnot0); case/card_gt0P => b _.
apply/card_gt0P.
have [tmp Htmp] : [finType of {f : {ffun A -> {ffun B -> 'I_n.+1}} |
                     \sum_(a1 in A) \sum_(b1 in B) f a1 b1 == n}].
  exists [ffun a1 => [ffun b1 => if (a1, b1) == (a, b) then
    Ordinal (ltnSn n) else Ordinal (ltn0Sn n)]].
  rewrite pair_big /=.
  rewrite (bigD1 (a, b)) //= big1 ; first by rewrite 2!ffunE eqxx addn0.
  move=> p /negbTE Hp ; by rewrite 2!ffunE -surjective_pairing Hp.
have Htmp' : (forall a b,
        (chan_of_jtype Anot0 Bnot0 tmp) a b =
        (let ln := \sum_(b0 in B) (tmp a) b0 in
         if ln == 0 then / #|B|%:R else (((tmp a) b)%:R / ln%:R)%R)).
  by move=> a0 b0; rewrite ffunE.
exists (@JType.mk _ _ _ (chan_of_jtype Anot0 Bnot0 tmp) tmp Htmp Htmp').
by rewrite inE.
Qed.

End jtype_facts.

Local Open Scope num_occ_scope.

Section shell_def.

Variables A B : finType.
Variable n : nat.
Variable ta : n.-tuple A.
Variable V : P_ n (A , B).

Local Open Scope nat_scope.

Definition shell :=
  [set tb : n.-tuple B | [forall a, [forall b, N(a, b |ta, tb) == (JType.f V) a b]]].

End shell_def.

Notation "V '.-shell' ta" := (shell ta V) : types_scope.

Section shelled_tuples_facts.

Variables A B : finType.
Variable n' : nat.
Let n := n'.+1.
Variable V : P_ n ( A , B).
Variable ta : n.-tuple A.
Variable tb : n.-tuple B.
Hypothesis Htb : tb \in V.-shell ta.

Lemma occ_co_occ : forall a b, N(a, b| ta, tb) = (JType.f V) a b.
Proof.
move: Htb => Htb' x0 y0 ; rewrite /shell inE in Htb'.
apply/eqP; move: y0; apply/forallP; move: x0; apply/forallP; exact Htb'.
Qed.

Variable P : P_ n ( A ).
Hypothesis Hta : ta \in T_{P}.

Lemma type_co_occ a b :
  ((type.f P) a * (JType.f V) a b = N(a, b | ta, tb) * N(a | ta))%nat.
Proof. by rewrite occ_co_occ mulnC (type_numocc Hta a). Qed.

End shelled_tuples_facts.

Section shelled_tuples_perm_facts.

Variables A B : finType.
Variable n : nat.
Variable V : P_ n ( A , B ).
Variable ta : n.-tuple A.
Variable s : 'S_n.

(* NB: Parameter imset : forall aT rT : finType, (aT -> rT) -> mem_pred aT -> {set rT}. *)
Lemma perm_Stuples_Stuples_perm : V.-shell (perm_tuple s ta) = (perm_tuple s) @: (V.-shell ta).
Proof.
rewrite /shell.
apply/eqP.
rewrite eqEsubset.
apply/andP; split; apply/subsetP => y.
- rewrite in_set => H.
  apply/imsetP.
  exists (perm_tuple (s^-1) y); last first.
    by rewrite perm_tuple_comp mulgV perm_tuple_id.
  rewrite in_set.
  apply/forallP => a; move/forallP/(_ a) in H.
  apply/forallP => b; move/forallP/(_ b) in H.
  by rewrite -(num_co_occ_perm _ _ _ _ s) perm_tuple_comp mulgV perm_tuple_id.
- case/imsetP => z.
  rewrite !in_set => H Hy.
  apply/forallP => a; move/forallP/(_ a) in H.
  apply/forallP => b; move/forallP/(_ b) in H.
  subst y.
  by rewrite num_co_occ_perm.
Qed.

Lemma card_shelled_tuples_perm : #| V.-shell ta | = #| V.-shell (perm_tuple s ta) |.
Proof.
rewrite perm_Stuples_Stuples_perm.
symmetry.
apply card_imset.
apply: perm_tuple_inj.
Qed.

End shelled_tuples_perm_facts.

Section sum_num_occ_tuple.

Variables A B : finType.
Variable n : nat.
Variable ta : n.-tuple A.
Variable k : 'I_#|A|.

Local Open Scope nat_scope.

Lemma sum_num_occ_size : forall (a : (sum_num_occ ta k.+1).-tuple B),
  sum_num_occ ta k <= size a.
Proof.
case=> a.
rewrite sum_num_occ_rec /= => /eqP ->.
by rewrite -{1}(addn0 (sum_num_occ ta k)) leq_add2l.
Qed.

Definition cons_tuples : (sum_num_occ ta k).-tuple B * N(enum_val k | ta).-tuple B ->
  (sum_num_occ ta k.+1).-tuple B.
case => a b.
exists [tuple of tval a ++ tval b].
by rewrite size_cat 2!size_tuple sum_num_occ_rec.
Defined.

Definition split_tuple : (sum_num_occ ta k.+1).-tuple B ->
  (sum_num_occ ta k).-tuple B * N(enum_val k | ta).-tuple B.
case => a b.
apply pair.
- apply: (@Tuple _ _ (take (sum_num_occ ta k) a)).
  rewrite size_take.
  move/eqP : b => ->.
  move: (sum_num_occ_inc_loc ta k); rewrite leq_eqVlt; case/orP.
  move/eqP => <-; by rewrite ltnn.
  by move=> ->.
- apply: (@Tuple _ _ (take N(enum_val k | ta) (drop (sum_num_occ ta k) a))).
  rewrite size_take size_drop.
  move/eqP : b => ->; by rewrite sum_num_occ_rec addKn ltnn.
Defined.

Lemma split_cons_tuples (Z : {set (sum_num_occ ta k.+1).-tuple B})
  (Z' : {set ((sum_num_occ ta k).-tuple B) * N(enum_val k | ta).-tuple B}) :
  forall tb, tb \in Z -> split_tuple tb \in Z' -> tb \in cons_tuples @: Z'.
Proof.
case=> sb Hsb sb_Z sb_Z'.
apply/imsetP.
exists (split_tuple (Tuple Hsb)) => //.
apply val_inj => /=.
rewrite -{1}(cat_take_drop (sum_num_occ ta k) sb).
f_equal.
rewrite take_oversize // size_drop.
clear -Hsb; move/eqP : Hsb => ->.
by rewrite sum_num_occ_rec addKn.
Qed.

Lemma subset_leq_card_split_tuple (Z : {set (sum_num_occ ta k.+1).-tuple B})
  (Z' : {set (sum_num_occ ta k).-tuple B * N(enum_val k | ta).-tuple B}) :
  (forall tb, tb \in Z -> split_tuple tb \in Z') ->
  #|Z| <= #|Z'|.
Proof.
move=> Hincl.
have <- : #| cons_tuples @: Z' | = #|Z'|.
  apply: card_in_imset; case=> [a1 a2] [b1 b2] Ha Hb [] /eqP.
  rewrite eqseq_cat; last by rewrite 2!size_tuple.
  case/andP => /eqP H1 /eqP H2.
  f_equal; by apply val_inj.
apply/subset_leq_card/subsetP => tb Htb.
by apply (split_cons_tuples Htb), Hincl.
Qed.

Hypothesis ta_sorted : sorted (@le_rank _) ta.

Local Open Scope tuple_ext_scope.

Lemma drop_take_is_filter : drop (sum_num_occ ta k) (take (sum_num_occ ta k.+1) ta) =
  filter (pred1 (enum_val k)) ta.
Proof.
apply (@eq_from_nth _ (enum_val k)) => [|i Hi].
  rewrite size_drop size_take size_tuple -/(minn (sum_num_occ ta k.+1) n) minn_sum_num_occ_n.
  by rewrite (sum_num_occ_sub _ k) /num_occ -size_filter.
rewrite size_drop size_take size_tuple -/(minn (sum_num_occ ta k.+1) n) minn_sum_num_occ_n ltn_subRL in Hi.
rewrite nth_drop nth_take => //.
have Hsumk : sum_num_occ ta k + i < n by apply (leq_trans Hi (sum_num_occ_leq_n ta k.+1)).
transitivity (enum_val k).
  transitivity (ta\_(Ordinal Hsumk)).
    by rewrite [in X in _ = X](tnth_nth (enum_val k)).
  apply sum_num_occ_enum_val => //=; by rewrite Hi andbT leq_addr.
rewrite -ltn_subRL (sum_num_occ_sub ta k) in Hi.
by rewrite filter_pred1_num_occ nth_nseq Hi.
Qed.

Hypothesis Bnot0 : (0 < #|B|)%nat.

Lemma drop_take_is_filter_zip (tb : n.-tuple B):
  drop (sum_num_occ ta k) (take (sum_num_occ ta k.+1) (zip ta tb)) =
  filter (fun p => p.1 == enum_val k) (zip ta tb).
Proof.
set a := enum_val k.
case/card_gt0P : Bnot0 => b _.
rewrite -(mkseq_nth (a, b) (zip ta tb)) -map_take -map_drop filter_map size_zip 2!size_tuple minnn.
f_equal.
rewrite drop_take_iota; last by rewrite sum_num_occ_inc_loc sum_num_occ_leq_n.
apply eq_in_filter => /= i.
rewrite mem_iota leq0n add0n /= => Hi.
rewrite nth_zip /=; last by rewrite 2!size_tuple.
transitivity (ta\_(Ordinal Hi) == a); by [rewrite -sum_num_occ_is_enum_val | rewrite (tnth_nth a)].
Qed.

Local Close Scope tuple_ext_scope.

Lemma drop_take_is_unzip2_filter (tb : n.-tuple B) :
  drop (sum_num_occ ta k) (take (sum_num_occ ta k.+1) tb) =
  unzip2 (filter (fun ab => ab.1 == enum_val k) (zip ta tb)).
Proof.
transitivity (unzip2 (drop (sum_num_occ ta k) (take (sum_num_occ ta k.+1) (zip ta tb)))).
by rewrite /unzip2 map_drop map_take -/(unzip2 (zip ta tb)) unzip2_zip // 2!size_tuple.
by rewrite drop_take_is_filter_zip.
Qed.

Lemma num_occ_num_co_occ (b : B) (sb : seq B) (tb : n.-tuple B) :
  sb = drop (sum_num_occ ta k) (take (sum_num_occ ta k.+1) tb) ->
  N(b | sb) = N(enum_val k, b | ta, tb).
Proof.
move=> Hsb.
set a := enum_val k.
rewrite /num_co_occ /num_occ // Hsb -!size_filter.
rewrite -(_ : filter (fun p => p.2 == b) (filter (fun p => p.1 == a) (zip ta tb)) =
              filter (pred1 (a, b)) (zip ta tb)) ; last first.
  rewrite -filter_predI /predI; apply eq_in_filter => i Hi /=.
  symmetry; by rewrite andbC {1}(surjective_pairing i) xpair_eqE.
set l1 := drop _ _.
set l2 := filter _ (zip ta tb).
suff -> : l1 = unzip2 l2 by rewrite filter_map size_map.
by rewrite /l1 /l2 {l1 l2} drop_take_is_unzip2_filter.
Qed.

End sum_num_occ_tuple.

Section take_shell_def.

Variables A B : finType.
Variable n : nat.
Variable ta : n.-tuple A.
Variable V : P_ n ( A , B).

Definition take_shell (k : nat) : {set (sum_num_occ ta k).-tuple B} :=
  (fun b : n.-tuple B =>
     tcast (minn_sum_num_occ_n ta k) [tuple of take (sum_num_occ ta k) b])
  @: (V.-shell ta).

(* Same set modulo cast: *)
Lemma full_take_shell : #| take_shell #|A| | = #| V.-shell ta |.
Proof.
apply card_imset; rewrite /injective => /= v v' vv'.
exact: (tcast_take_inj (full_sum_num_occ_n ta (leqnn #|A|)) vv').
Qed.

End take_shell_def.

Section row_num_occ_sect.

Variables (A : finType) (n : nat).
Variable P : P_ n ( A ).
Variable B : finType.
Variable V : P_ n (A , B).

Definition row_num_occ := forall ta, ta \in T_{P} ->
  forall a, (\sum_(b in B) (JType.f V) a b)%nat = N(a | ta).

Hypothesis H : row_num_occ.
Variable ta : n.-tuple A.
Hypothesis Hta : ta \in T_{P}.
Variable a : A.
Variable b : B.

Lemma ctyp_element_ub : ((JType.f V) a b < N(a | ta).+1)%nat .
Proof. by rewrite ltnS -(H Hta) (bigD1 b) // leq_addr. Qed.

End row_num_occ_sect.

Section take_shell_row_num_occ.

Variables A B : finType.
Variable n : nat.
Variable V : P_ n ( A , B).
Variable P : P_ n ( A ).
Variable ta : n.-tuple A.
Hypothesis Hta : ta \in T_{P}.
Hypothesis Hrow_num_occ : row_num_occ P V.

Local Open Scope nat_scope.

Definition type_of_row (a : A) (Ha : N(a | ta) != 0) : P_ N(a | ta) ( B ).
pose f := [ffun b => Ordinal (ctyp_element_ub Hrow_num_occ Hta a b)].
pose d := [ffun b => ((f b)%:R / N(a | ta)%:R)%R].
have d0 : forall b, (0 <= d b)%R.
  move=> b.
  rewrite /d /= ffunE.
  apply mulR_ge0; first exact/leR0n.
  apply/invR_ge0/ltR0n; by rewrite lt0n.
have d1 : (\sum_(b : B) d b)%R = 1%R.
  under eq_bigr do rewrite ffunE /=.
  rewrite -big_distrl /= -big_morph_natRD.
  set lhs := \sum_i _.
  suff -> : lhs = N(a | ta) by rewrite mulRV // INR_eq0'.
  rewrite /lhs /f /= -[in X in _ = X](Hrow_num_occ Hta a).
  apply eq_bigr => b _; by rewrite ffunE.
by apply (@type.mkType _ _ (FDist.make d0 d1) f) => b; rewrite ffunE.
Defined.

Hypothesis ta_sorted : sorted (@le_rank _) ta.
Hypothesis Bnot0 : (0 < #|B|)%nat.

Lemma card_take_shell_incl0 (k : 'I_#|A|) (Ha : 0 = N(enum_val k | ta)) tb :
  tb \in take_shell ta V k.+1 ->
  split_tuple tb \in setX (take_shell ta V k) (set1 (tcast Ha [tuple of ([::] : seq B)])).
Proof.
move=> tbV.
destruct tb as [s Hs].
move: tbV; rewrite /take_shell [in X in X -> _]/= => /imsetP[tb' tb'V].
move/esym/tcast2tval; rewrite [X in X -> _]/= => tb's.
rewrite /split_tuple /= in_setX; apply/andP; split.
- apply/imsetP; exists tb' => //.
  apply/val_inj => /=.
  rewrite eq_tcast /=.
  by rewrite -tb's sum_num_occ_rec take_take // leq_addr.
- rewrite in_set.
  apply/eqP/val_inj => /=.
  by rewrite eq_tcast -Ha take0.
Qed.

Lemma card_take_shell_incl (k : 'I_#|A|) (Ha : N(enum_val k | ta) != 0 ) tb :
  tb \in take_shell ta V k.+1 ->
  split_tuple tb \in setX (take_shell ta V k) (T_{type_of_row Ha}).
Proof.
move=> Hlhs.
destruct tb as [sb Hsb].
rewrite /take_shell /= in Hlhs. move:Hlhs.
case/imsetP => tb Htb_1 Htb_2.
symmetry in Htb_2; move/tcast2tval in Htb_2; rewrite /= in Htb_2.
rewrite /split_tuple /= in_setX.
apply/andP; split.
- apply/imsetP; exists tb => //; apply/val_inj => /=.
  by rewrite eq_tcast /= -Htb_2 sum_num_occ_rec take_take // leq_addr.
- rewrite in_set.
  set t := Tuple _.
  have Ht : tval t = take N(enum_val k | ta) (drop (sum_num_occ ta k) sb) by [].
  apply/forallP => b.
  rewrite /= ffunE /=.
  move: Htb_1.
  rewrite in_set ffunE.
  move/forallP/(_ (enum_val k))/forallP/(_ b)/eqP => /= H.
  rewrite -{1}H.
  rewrite -Htb_2 in Ht.
  apply/eqP.
  have Ht2 : tval t = drop (sum_num_occ ta k) (take (sum_num_occ ta k.+1) tb).
    rewrite Ht {1}sum_num_occ_rec take_drop take_take; last by rewrite addnC.
    by rewrite addnC sum_num_occ_rec.
  congr (_ %:R / _%:R)%R.
  exact/esym/num_occ_num_co_occ.
Qed.

Lemma card_take_shell (k : 'I_#|A|) (Ha : N(enum_val k | ta) != 0) :
  #|take_shell ta V k.+1| <= #|take_shell ta V k| * #| T_{type_of_row Ha} |.
Proof.
rewrite -cardX -[X in _ <= X]cardsE.
apply (subset_leq_card_split_tuple (card_take_shell_incl Ha)).
Qed.

Lemma card_take_shell0 (k : 'I_#|A|) (Ha : N(enum_val k | ta) == 0) :
  #|take_shell ta V k.+1| <= #|take_shell ta V k|.
Proof.
move/eqP : (Ha).
move/esym => Ha'.
eapply leq_trans; first by apply (subset_leq_card_split_tuple (card_take_shell_incl0 Ha')).
by rewrite cardsX cards1 muln1.
Qed.

Definition card_type_of_row (i : 'I_#|A|) :=
  match Bool.bool_dec (N(enum_val i | ta) == O) true with
  | left _ => 1
  | right Ha => #| T_{type_of_row (Bool.eq_true_not_negb _ Ha)} |
  end.

Lemma split_nocc_rec : forall k, k <= #|A| ->
  #|take_shell ta V k| <= \prod_ (i < #|A| | i < k) card_type_of_row i.
Proof.
elim.
- move=> _; rewrite big_pred0; last by move=> ? /=; rewrite ltn0.
  rewrite -(expn0 #|B|) -[X in _ <= expn _ X](sum_num_occ_0 ta) -card_tuple -cardsT.
  by apply subset_leq_card, subsetT.
- move=> k IH HSk /=.
  move: (IH (ltnW HSk)) => {}IH.
  rewrite (bigD1 (Ordinal HSk)) //=.
  rewrite (eq_bigl (fun i : 'I_#|A| => i < k) _); last first.
    move=> i /=.
    case/boolP : (i < k) => Hcase.
    - have -> : i != Ordinal HSk by rewrite neq_ltn; apply/orP; apply or_introl.
      by rewrite andbC /= ltnW.
    - rewrite andbC -ltn_neqAle.
      by move/negbTE : Hcase => ->.
  rewrite /card_type_of_row; case: Bool.bool_dec => [e|/Bool.eq_true_not_negb e].
    rewrite mul1n.
    eapply leq_trans; [exact: (card_take_shell0 e) | by []].
  apply (leq_trans (card_take_shell e)).
  rewrite mulnC leq_pmul2l //.
  apply/card_gt0P.
  set Q := type_of_row e.
  case: (typed_tuples_not_empty_alt e Q) => tb Htb.
  by exists tb.
Qed.

Lemma card_shelled_tuples_leq_prod_card : #| V.-shell ta | <= \prod_ ( i < #|A|) card_type_of_row i.
Proof.
rewrite -full_take_shell [X in _ <= X](_ : _ = \prod_(i < #|A| | i < #|A|) card_type_of_row i); last first.
  apply eq_bigl => ?; symmetry; by apply ltn_ord.
exact (split_nocc_rec (leqnn #|A|)).
Qed.

End take_shell_row_num_occ.

Local Open Scope entropy_scope.

Section card_shell_ub.
Variables (A B : finType) (n' : nat).
Let n := n'.+1.
Variable V : P_ n ( A , B).
Variable P : P_ n ( A ).
Variable ta : n.-tuple A.
Hypothesis Hta : ta \in T_{P}.
Hypothesis Vctyp : row_num_occ P V.
Hypothesis ta_sorted : sorted (@le_rank _) ta.
Hypothesis Bnot0 : (0 < #|B|)%nat.

Lemma card_shell_leq_exp_entropy :
  #| V.-shell ta |%:R <= exp2 (n%:R * `H(V | P)).
Proof.
rewrite cond_entropy_chanE2.
apply (@leR_trans (\prod_ ( i < #|A|) card_type_of_row Hta Vctyp i)%:R).
- exact/le_INR/leP/card_shelled_tuples_leq_prod_card.
- rewrite exp2_pow big_morph_natRM.
  rewrite (@big_morph _ _ (fun r : R => ((exp2 r) ^ n)%R) 1%R Rmult _ Rplus _); last 2 first.
    move=> a b /=; rewrite -!exp2_pow mulRDr /exp2 !ExpD; by field.
    by rewrite -exp2_pow mulR0 /exp2 Exp_0.
  rewrite (reindex_onto (fun x => enum_rank x) (fun y => enum_val y)) => [|i _]; last by rewrite enum_valK.
  rewrite (_ : \prod_(j | enum_val (enum_rank j) == j) _ =
               \prod_(j : A) (card_type_of_row Hta Vctyp (enum_rank j))%:R); last first.
      apply eq_bigl => a; rewrite enum_rankK; by apply/eqP.
  apply leR_prodR => a.
  split; first exact/leR0n.
  rewrite -exp2_pow mulRA.
  rewrite /card_type_of_row; case: Bool.bool_dec => [e|/Bool.eq_true_not_negb e].
    rewrite -[X in X <= _]exp2_0.
    apply Exp_le_increasing, mulR_ge0 => //.
      apply mulR_ge0 => //; exact: leR0n.
      exact: entropy_ge0.
  set pta0 := type_of_row Hta Vctyp _.
  rewrite (_ : exp2 _ = exp2 (N(a | ta)%:R * `H pta0)%R).
    by rewrite -[in X in _ <= exp2 (X * _)](enum_rankK a); apply card_typed_tuples.
  congr (exp2 (_ * _)).
  + by rewrite -type_fun_type // (type_numocc Hta).
  + rewrite /entropy.
    apply Ropp_eq_compat, eq_bigr => b _.
    rewrite /pta0 (JType.c_f V) /= (Vctyp Hta a) -{1 4}(enum_rankK a).
    by move/negbTE : (e) => ->; rewrite !ffunE /= enum_rankK.
Qed.

End card_shell_ub.

(* TODO: move? *)
Lemma map_pred1_nseq {A : eqType} : forall (l : seq A) n a, a \notin l ->
  map (pred1 a) (flatten [seq nseq (n x) x | x <- l]) = nseq (\sum_(i <- l) (n i)) false.
Proof.
elim.
  move=> n0 a Ha /=;  by rewrite big_nil.
move=> h t IH n0 a.
rewrite in_cons negb_or.
case/andP => H1 H2 /=.
rewrite map_cat IH // (_ : map _ _ = nseq (n0 h) false); last first.
  by rewrite map_nseq /= -(negbTE H1) eqtype.eq_sym.
by rewrite big_cons nseq_add.
Qed.

(* TODO: move? *)
Lemma map_filter_nseq_nil {A : eqType} : forall (l : seq A) (p : pred A) n,
  (forall x, x \in l -> ~~ p x) ->
  flatten (map (filter p) (map (fun x => nseq (n x) x) l)) = [::].
Proof.
elim=> // h t IH p n0 H /=.
rewrite IH.
- rewrite cats0.
  transitivity (filter xpred0 (nseq (n0 h) h)).
    apply eq_in_filter => a /nseqP[-> _].
    apply/negP/negP/H; by rewrite in_cons eqxx.
  by rewrite filter_pred0.
move=> x0 Hx0.
apply H; by rewrite in_cons Hx0 orbC.
Qed.

(* TODO: move? *)
Lemma map_filter_pred1_nseq {A : eqType} (a : A) : forall lst n, uniq lst -> a \in lst ->
  flatten (map (filter (pred1 a)) (map (fun x => nseq (n x) x) lst)) = nseq (n a) a.
Proof.
elim=> // h t IH n0 /=.
case/andP=> H1 H2.
rewrite in_cons.
case/orP.
  move/eqP => ?; subst a.
  rewrite {IH} map_filter_nseq_nil //; last first.
    move=> x0 x0_t /=.
    apply/eqP => ?; subst x0.
    by rewrite x0_t in H1.
  rewrite (_ : filter _ _ = filter xpredT (nseq (n0 h) h)); last first.
    apply eq_in_filter => i /nseqP[-> /= ]; by rewrite eqxx.
  by rewrite filter_predT /= cats0.
move=> a_t.
rewrite IH // (_ : filter _ _ = filter xpred0 (nseq (n0 h) h)); last first.
  apply eq_in_filter => i /nseqP[-> /= _].
  by apply: contraNF H1 => /eqP  ->.
by rewrite filter_pred0.
Qed.

Section shell_not_empty_sorted.

Variables A B : finType.
Variable n : nat.
Variable ta : n.-tuple A.
Hypothesis ta_sorted : sorted (@le_rank A) ta.
Variable V : P_ n ( A , B ).
Variable P : P_ n ( A ).
Hypothesis Hrow_num_occ : row_num_occ P V.
Hypothesis Hta : ta \in T_{P}.

Local Open Scope nat_scope.

Lemma shell_not_empty' : exists sb : seq B,
  exists Hsb : size sb == n, Tuple Hsb \in V.-shell ta.
Proof.
exists (flatten (map (fun a => flatten (map (fun b => nseq (JType.f V a b) b) (enum B))) (enum A))).
set cdom := flatten _.
have /eqP Hy : size cdom = n.
  rewrite size_flatten /shape -map_comp sumn_big_addn big_map [LHS]big_filter.
  rewrite -[RHS](sum_num_occ_all ta).
  rewrite [RHS](reindex_onto enum_rank enum_val) => [|i _] ; last by rewrite enum_valK.
  apply eq_big.
  - move=> ? /=; by rewrite enum_rankK eqxx.
  - move=> a _ /=.
    rewrite [in RHS]enum_rankK.
    rewrite size_flatten /shape -map_comp sumn_big_addn big_map big_filter.
    transitivity (\sum_(b : B) JType.f V a b).
      apply eq_bigr => ? _ /=; by rewrite size_nseq.
    by rewrite (Hrow_num_occ Hta).
exists Hy.
rewrite in_set.
apply/forallP => a.
apply/forallP => b.
rewrite /num_co_occ /num_occ /=.
apply/eqP.
rewrite (@sorted_is_flattened _ _ (@le_rank_trans _)
  (@le_rank_asym _) (@le_rank_refl _) #|A| (enum A) ta) //; last 4 first.
  by rewrite cardE.
  by apply enum_uniq.
  by apply sorted_enum.
  move=> ? _; by rewrite mem_enum.
have -> : [seq filter (pred1 elt) ta | elt <- enum A] = [seq nseq N(x0 | ta) x0 | x0 <- enum A].
  apply eq_in_map => ? _; by rewrite filter_pred1_num_occ.
have sz_flat : forall a, size (flatten [seq nseq (JType.f V a b) b | b <- enum B]) = N(a | ta).
  move=> a'; rewrite size_flatten /shape -map_comp sumn_big_addn big_map.
  rewrite -(Hrow_num_occ Hta a') /= enumT.
  apply eq_bigr => b' _; by rewrite size_nseq.
set dom := flatten _.
have Hdom : size dom = n.
  rewrite /dom size_flatten /shape -map_comp sumn_big_addn big_map.
  rewrite big_filter -[RHS](sum_num_occ_alt ta).
  apply eq_bigr => a' _ /=; by rewrite size_nseq.
move/eqP in Hy.
rewrite -size_filter (_ : filter _ _ = filter
  (predI (fun x => x.2 == b) (fun x => x.1 == a)) (zip dom cdom)); last first.
  apply eq_in_filter; case=> i1 i2 Hi /=; by rewrite xpair_eqE andbC.
rewrite filter_predI (@filter_zip_L _ _ n) //.
have -> : mask (map (pred1 a) dom) cdom = flatten [seq nseq (JType.f V a b) b | b <- enum B].
  have [A1 [A2 A12]] : exists A1 A2, enum A = A1 ++ a :: A2.
    have /splitPr[A1 A2] : a \in enum A by rewrite mem_enum.
    by exists A1, A2.
  rewrite /cdom /dom A12 map_cat flatten_cat map_cat.
  rewrite [in X in mask _ X]map_cat flatten_cat mask_cat; last first.
    rewrite size_map size_flatten /shape -map_comp sumn_big_addn big_map.
    rewrite size_flatten /shape -map_comp sumn_big_addn big_map.
    apply eq_bigr => i _ /=; by rewrite sz_flat size_nseq.
  rewrite (_ : _ :: _ = [:: a] ++ A2) //.
  rewrite map_cat.
  rewrite [in X in _ ++ mask _ X = _]map_cat flatten_cat.
  rewrite mask_cat; last first.
    by rewrite size_map /= cats0 sz_flat size_nseq.
  transitivity (mask (map (pred1 a) (flatten [seq nseq N(a1 | ta) a1 | a1 <- [:: a]]))
     (flatten [seq flatten [seq nseq (JType.f V a1 b1) b1 | b1 <- enum B] | a1 <- [:: a]])).
    move: (enum_uniq A); rewrite A12 cat_uniq /= negb_or /=.
    case/andP => uniqA1 /andP [] /andP [a_A1 _] /andP [a_A2 _].
    rewrite {1}(_ : mask _ _ = [::]); last by rewrite map_pred1_nseq // mask_false.
    rewrite (_ : mask (map _ (flatten [seq _ a | a <- A2])) _ = [::]); last first.
      by rewrite map_pred1_nseq // mask_false.
    by rewrite 3!cats0 cat0s.
  rewrite /= 2!cats0 map_nseq /= eqxx mask_true // sz_flat; by apply eq_leq.
rewrite (@filter_zip_R _ _ N(a | ta)) //; last first.
  rewrite /dom filter_flatten size_flatten /shape -!map_comp sumn_big_addn.
  rewrite big_map (bigD1 a) // big1 /= => [|a1 Ha1].
  - rewrite (_ : filter _ _ = filter (xpredT) (nseq N(a | ta) a)); last first.
      apply eq_in_filter => i /nseqP[-> _]; by rewrite /pred1 /= eqxx.
    by rewrite addn0 filter_predT size_nseq.
  - rewrite (_ : filter _ _ = filter (xpred0) (nseq N(a1 | ta) a1)); last first.
      apply eq_in_filter => j /nseqP[-> _].
      by rewrite /pred1 /= (negbTE Ha1).
    by rewrite filter_pred0.
rewrite size_zip size_mask; last first.
  rewrite size_map /dom sz_flat (_ : filter _ _ = nseq N(a | ta) a); last first.
    by rewrite filter_flatten map_filter_pred1_nseq // ?enum_uniq // ?mem_enum.
  by rewrite size_nseq.
rewrite size_filter.
set x1 := count _ _. set x2 := count _ _.
have -> : x1 = x2 by rewrite /x1 /x2 count_map; apply eq_in_count.
have -> : x2 = JType.f V a b.
  rewrite /x2 -size_filter (_ : filter _ _ = nseq (JType.f V a b) b); last first.
    by rewrite filter_flatten map_filter_pred1_nseq // ?enum_uniq // ?mem_enum.
  by rewrite size_nseq.
by apply/minn_idPl.
Qed.

Lemma shell_not_empty_sorted : exists tb, tb \in V.-shell ta.
Proof. case: (shell_not_empty') => tb [Htb H]; by exists (Tuple Htb). Qed.

End shell_not_empty_sorted.

Section shell_not_empty.

Variables A B : finType.
Variable n : nat.
Variable ta : n.-tuple A.
Variable V : P_ n ( A , B ).
Variable P : P_ n ( A ).
Hypothesis Hrow_num_occ : row_num_occ P V.
Hypothesis Hta : ta \in T_{P}.

Local Open Scope nat_scope.

Lemma shell_not_empty : exists tb, tb \in V.-shell ta.
Proof.
case: (tuple_exist_perm_sort (@le_rank A) ta) => s /=.
rewrite /sort_tuple /=.
set t' := Tuple _.
move=> ta_t'.
have t'ta : t' = perm_tuple s^-1 ta.
  by rewrite ta_t' perm_tuple_comp mulVg perm_tuple_id.
have sorted_t' : sorted (@le_rank A) t' by exact/sort_sorted/le_rank_total.
have Ht' : t' \in T_{P} by rewrite t'ta; apply/perm_tuple_in_Ttuples.
case: (shell_not_empty_sorted sorted_t' Hrow_num_occ Ht') => tb Htb.
by exists (perm_tuple s tb); rewrite ta_t' perm_Stuples_Stuples_perm imset_f.
Qed.

End shell_not_empty.

Section cond_type_def.

Variable A : finType.
Variable n : nat.
Variable P : P_ n ( A ).
Variable B : finType.

Definition cond_type :=
  [set V : P_ n ( A , B ) | [forall ta, (ta \in T_{P}) ==> (V.-shell ta != set0)]].

End cond_type_def.

Notation "'\nu_' n '^{' A ',' B '}' '(' P ')'" :=
  (@cond_type A n P B) : types_scope.
Notation "'\nu^{' B '}' '(' P ')'" := (@cond_type _ _ P B) : types_scope.

Section cond_type_prop.

Variable A : finType.
Variable n : nat.
Variable P : P_ n ( A ).
Variable B : finType.

Local Open Scope nat_scope.

Lemma card_nu : #|\nu^{B}( P )| <= expn n.+1 (#|A| * #|B|).
Proof.
apply: (leq_trans _ (bound_card_jtype A B n)).
apply subset_leq_card; by apply/subsetP.
Qed.

Lemma shell_injective (V V' : P_ n (A , B)) (Vctyp : V \in \nu^{B}(P))
  ta (Hta : ta \in T_{P}) : V.-shell ta = V'.-shell ta ->
  forall a b, (JType.f V) a b = (JType.f V') a b.
Proof.
move=> H a b.
rewrite /cond_type in_set in Vctyp.
move/forallP in Vctyp.
move: {Vctyp}(Vctyp ta) => Vctyp.
move/implyP in Vctyp.
move: {Vctyp}(Vctyp Hta) => Vctyp.
case/set0Pn : Vctyp => tb V_tb.
have V'_tb : tb \in V'.-shell ta by rewrite -H.
apply val_inj => /=.
move: V_tb; rewrite in_set => /forallP/(_ a)/forallP/(_ b)/eqP => <-.
move: V'_tb; by rewrite in_set => /forallP/(_ a)/forallP/(_ b)/eqP => <-.
Qed.

End cond_type_prop.

Section cond_type_equiv_sect.

Variable A : finType.
Variable n : nat.
Variable P : P_ n ( A ).
Variable B : finType.
Variable V : P_ n ( A , B ).

Lemma cond_type_equiv : row_num_occ P V <-> [forall ta, (ta \in T_{P}) ==> (V.-shell ta != set0)].
Proof.
split=> H.
- apply/forallP => ta.
  apply/implyP => Hta.
  apply/set0Pn.
  by apply shell_not_empty with P.
- move=> ta Hta a.
  move/forallP : H => /(_ ta)/implyP/(_ Hta).
  case/set0Pn => tb.
  rewrite /shell in_set.
  move/forallP/(_ a)/forallP => H.
  rewrite -(num_co_occ_num_occ ta tb).
  apply eq_bigr => b _.
  by move/eqP : (H b) => <-.
Qed.

End cond_type_equiv_sect.

Module OutType.

Section OutType_sect.

Local Open Scope nat_scope.

Variables A B : finType.
Variable n' : nat.
Let n := n'.+1.
Variable V : P_ n ( A , B ).

Definition f := [ffun b => ((\sum_(a in A) (JType.f V) a b)%:R / n%:R)%R].

Lemma f0 (b : B) : (0 <= f b)%R.
Proof. rewrite ffunE; apply divR_ge0; [exact/leR0n | exact/ltR0n]. Qed.

Lemma f1 : (\sum_(b in B) f b = 1)%R.
Proof.
under eq_bigr do rewrite ffunE /=.
rewrite -big_distrl /= -big_morph_natRD exchange_big /=.
by move/eqP : (JType.sum_f V) => ->; rewrite mulRV // INR_eq0'.
Qed.

Definition d : fdist B := FDist.make f0 f1.

Definition P : P_ n ( B ).
refine (@type.mkType _ _ (FDist.make f0 f1) [ffun b => Ordinal (jtype_entry_ub V b)] _).
by move=> b /=; rewrite !ffunE.
Defined.

End OutType_sect.

End OutType.

Notation "'`tO(' V )" := (OutType.P V) : types_scope.

Section output_type_facts.

Variables A B : finType.
Variable n' : nat.
Let n := n'.+1.
Variable V : P_ n ( A , B ).
Variable P : P_ n ( A ).

Lemma shell_subset_output_type ta : V.-shell ta \subset T_{ `tO( V ) }.
Proof.
apply/subsetP => tb; rewrite 2!in_set => Htb.
apply/forallP => b.
apply/eqP.
rewrite /OutType.P /OutType.d /OutType.f ffunE.
do 2 f_equal.
rewrite -(num_co_occ_partial_sum_alt tb ta).
apply eq_bigr => a _.
rewrite num_co_occ_sym.
by move/forallP/(_ a)/forallP/(_ b)/eqP : Htb.
Qed.

Hypothesis Bnot0 : (0 < #|B|)%nat.
Hypothesis Vctyp : V \in \nu^{B}(P).

Lemma output_type_out_fdist : forall b, (`tO( V )) b = `O( P , V ) b.
Proof.
rewrite /fdist_of_ffun /= /OutType.d /OutType.f => b /=.
rewrite ffunE big_morph_natRD /Rdiv (big_morph _ (morph_mulRDl _) (mul0R _)).
rewrite fdist_outE; apply eq_bigr => a _.
case: (typed_tuples_not_empty P) => /= ta Hta.
move: (Vctyp).
rewrite in_set.
move/cond_type_equiv/(_ _ Hta a) => sum_V.
rewrite (JType.c_f V) /=.
case: ifP => [/eqP |] Hcase.
  rewrite Hcase in sum_V.
  rewrite in_set in Hta.
  move/forallP/(_ a) : Hta.
  rewrite -sum_V div0R.
  move/eqP => ->; rewrite mulR0.
  move/eqP in Hcase.
  rewrite sum_nat_eq0 in Hcase.
  move/forallP/(_ b) : Hcase.
  move/implyP/(_ Logic.eq_refl)/eqP => ->.
  by rewrite mul0R.
- rewrite -mulRA sum_V; congr (_ * _).
  move: Hta; rewrite in_set => /forallP/(_ a)/eqP ->.
  by rewrite mulRA -{1}(mul1R (/ n%:R)) mulVR // INR_eq0' -sum_V Hcase.
Qed.

Lemma output_type_out_entropy : `H (`tO( V )) = `H(P `o V).
Proof.
rewrite /entropy; f_equal.
apply eq_bigr => b _; by rewrite output_type_out_fdist.
Qed.

End output_type_facts.

Section card_perm_shell.

Variables A B : finType.
Variable n' : nat.
Let n := n'.+1.
Variable P : P_ n ( A ).
Variable ta : n.-tuple A.
Variable V : P_ n ( A , B ).
Hypothesis Hta : ta \in T_{P}.
Hypothesis Vctyp : V \in \nu^{B}(P).
Hypothesis Bnot0 : (0 < #|B|)%nat.

Lemma card_shelled_tuples : #| V.-shell ta |%:R <= exp2 (n%:R * `H(V | P)).
Proof.
case: (tuple_exist_perm_sort (@le_rank A) ta) => /= s Hta'.
have H : sort (@le_rank _) ta =
  perm_tuple (s^-1) ta by rewrite {2}Hta' perm_tuple_comp mulVg perm_tuple_id.
have {}Hta': ta = perm_tuple s (sort_tuple (@le_rank _) ta) by rewrite {1}Hta'.
rewrite (card_shelled_tuples_perm V ta (s^-1)).
rewrite Hta' perm_tuple_comp mulVg perm_tuple_id.
apply card_shell_leq_exp_entropy => //.
- rewrite in_set.
  apply/forallP => a /=.
  rewrite H num_occ_perm.
  move: a; apply/forallP.
  move: Hta; by rewrite in_set.
- apply cond_type_equiv.
  move: (Vctyp); by rewrite in_set.
- exact/sort_sorted/le_rank_total.
Qed.

End card_perm_shell.

Section shell_partition.

Local Open Scope fun_scope.
Local Open Scope nat_scope.

Variables A B : finType.
Variable n' : nat.
Let n := n'.+1.

(** The stochastic matrix with entries N(a, b | ta, tb): *)

Hypothesis Anot0 : (0 < #|A|)%nat.
Hypothesis Bnot0 : (0 < #|B|)%nat.

Definition num_co_occ_jtype (ta : n.-tuple A) (tb : n.-tuple B) : P_ n (A , B).
set f := [ffun a => [ffun b => Ordinal (num_co_occ_ub a b ta tb)]].
have Hf : \sum_(a in A) \sum_(b in B) f a b == n.
  rewrite /f.
  apply/eqP.
  transitivity (\sum_a \sum_b N(a, b|ta, tb)); last by rewrite num_co_occ_sum.
  apply eq_big => a //= _.
  apply eq_big => b //= _.
  by rewrite 2!ffunE.
have Htmp' : (forall a b,
        (chan_of_jtype Anot0 Bnot0 f) a b =
        (let ln := (\sum_(b0 in B) (f a) b0)%nat in
         if ln == O then / #|B|%:R else (((f a) b)%:R / ln%:R))%R).
  by move=> a b; rewrite ffunE.
exact (@JType.mk _ _ _ (chan_of_jtype Anot0 Bnot0 f) f Hf Htmp').
Defined.

Definition relYn (ta : n.-tuple A) (tb tb' : n.-tuple B) :=
  [forall a, [forall b, N(a, b| ta, tb) == N(a, b|ta, tb')] ].

Lemma reflexive_relYn ta : reflexive (relYn ta).
Proof.
rewrite /reflexive /relYn => tb.
apply/forallP => a; apply/forallP => b; by rewrite eqxx.
Qed.

Variable ta : n.-tuple A.
Variable P : P_ n ( A ).
Hypothesis Hta : ta \in T_{P}.

Definition shell_partition : {set set_of_finType [finType of n.-tuple B]} :=
  (fun V => V.-shell ta) @: \nu^{B}(P).

Lemma cover_shell : cover shell_partition = [set: n.-tuple B].
Proof.
rewrite /cover /cond_type.
apply/setP => tb.
rewrite in_set.
apply/bigcupP.
exists (num_co_occ_jtype ta tb).-shell ta.
- apply/imsetP; exists (num_co_occ_jtype ta tb) => //.
  rewrite inE.
  apply cond_type_equiv => ta' Hta' a.
  transitivity (\sum_(b in B) N(a, b | ta, tb)).
  - apply eq_bigr => b _.
    by rewrite /num_co_occ_jtype /= 2!ffunE.
  - rewrite num_co_occ_partial_sum_alt.
    move: Hta'; rewrite in_set => /forallP/(_ a)/eqP => Hta'.
    move: Hta.
    rewrite in_set => /forallP/(_ a)/eqP.
    rewrite Hta' eqR_mul2r; last by apply/invR_neq0; rewrite INR_eq0.
    by move/INR_eq.
- rewrite in_set.
  apply/forallP => a. apply/forallP => b.
  by rewrite /num_co_occ_jtype /= 2!ffunE.
Qed.

Lemma trivIset_shell' tb tb' V : tb \in V.-shell ta -> tb' \in V.-shell ta = relYn ta tb' tb.
Proof.
rewrite 2!in_set => H.
rewrite /relYn.
apply eq_forallb => a; apply eq_forallb => b; move: H.
by move/forallP/(_ a)/forallP/(_ b)/eqP => ->.
Qed.

Lemma trivIset_shell : trivIset shell_partition.
Proof.
apply/trivIsetP => S1 S2; case/imsetP => V1 _ HS1; case/imsetP => V2 _ HS2 HS12.
subst S1 S2.
rewrite /disjoint.
apply/pred0P => tb /=.
apply/negP/negP.
move: tb.
apply/forallP; rewrite -negb_exists; apply/negP; case/existsP => tb /andP [H1 H2]; contradict HS12.
apply/negP/negPn/eqP/setP => ?.
rewrite 2!(@trivIset_shell' tb) //.
apply reflexive_relYn.
Qed.

End shell_partition.

Section sum_tuples_ctypes.
Variables (A B : finType) (n' : nat).
Let n := n'.+1.
Variable ta : n.-tuple A.
Variable P : P_ n ( A ).
Hypothesis Hta : ta \in T_{P}.

Let sum_tuples_ctypes'' f :
  \sum_ (S | S \in shell_partition B ta P) \sum_(tb in S) f tb =
  \sum_ (V | V \in \nu^{B}(P)) \sum_ (tb in V.-shell ta) f tb.
Proof.
rewrite big_imset // => V V' HV HV' /=.
move/(shell_injective _) => /(_ P HV Hta) V_V' {HV HV'}.
case : V V_V' => /= c fV HfV Hc HVV'.
case : V' HVV' => /= c' fV' HfV' Hc' fV_fV'.
have ? : fV = fV' by apply/ffunP => a; apply/ffunP => b; rewrite fV_fV'.
subst fV.
by apply/jtype_eqP => /=.
Qed.

Hypothesis Anot0 : (0 < #|A|)%nat.
Hypothesis Bnot0 : (0 < #|B|)%nat.

Let sum_tuples_ctypes' f : \sum_ (tb : _ ) f tb =
  \sum_ (V | V \in \nu^{B}(P)) \sum_ (tb in V.-shell ta) f tb.
Proof.
transitivity (\sum_ (tb in [set: n.-tuple B]) f tb).
  by apply eq_bigl => tb; rewrite in_set.
rewrite -(cover_shell Anot0 Bnot0 Hta).
rewrite -sum_tuples_ctypes'' // big_trivIset //.
apply trivIset_shell.
Qed.

Lemma sum_tuples_ctypes f F :
  \sum_(tb | F tb) f tb =
  \sum_(V | V \in \nu^{B}(P)) \sum_ (tb in V.-shell ta | F tb) f tb.
Proof.
rewrite big_mkcond /=.
transitivity (\sum_(V | V \in \nu^{B}(P)) \sum_(tb in V.-shell ta) if F tb then f tb else 0).
  by apply sum_tuples_ctypes'.
apply eq_bigr => s _.
rewrite [in LHS]big_mkcond /= [in RHS]big_mkcond /=.
apply/esym/eq_bigr => tb _.
by case/boolP : (tb \in s.-shell ta).
Qed.

End sum_tuples_ctypes.


Require Import Coq.Arith.Compare_dec.
Require Export Fiat.Common.Coq__8_4__8_5__Compat.

Generalizable All Variables.

Class Decidable (P : Prop) := {
  Decidable_witness : bool;
  Decidable_spec : Decidable_witness = true <-> P
}.
Arguments Decidable_witness {P _} /.
Arguments Decidable_spec {P _} /.

Lemma Decidable_sound : forall P (H : Decidable P),
  Decidable_witness = true -> P.
Proof. intros; apply Decidable_spec; trivial. Qed.

Lemma Decidable_complete : forall P (H : Decidable P),
  P -> Decidable_witness = true.
Proof. intros; apply Decidable_spec; trivial. Qed.

Lemma Decidable_sound_alt : forall P (H : Decidable P),
  ~ P -> Decidable_witness = false.
Proof.
  intros.
  destruct Decidable_witness eqn:Heqe; trivial.
  apply Decidable_spec in Heqe.
  contradiction.
Qed.

Lemma Decidable_complete_alt : forall P (H : Decidable P),
  Decidable_witness = false -> ~ P.
Proof.
  intros.
  unfold not; intros.
  apply Decidable_spec in H1.
  congruence.
Qed.

Definition decide P {H : Decidable P} := Decidable_witness (Decidable:=H).

Ltac _decide_ P H :=
  let b := fresh "b" in
  set (b := decide P) in *;
  assert (H : decide P = b) by reflexivity;
  clearbody b;
  destruct b; [apply Decidable_sound in H|apply Decidable_complete_alt in H].

Tactic Notation "decide" constr(P) "as" ident(H) :=
  _decide_ P H.

Tactic Notation "decide" constr(P) :=
  let H := fresh "H" in _decide_ P H.

#[global]
Hint Extern 2 =>
  match goal with
    [ H : @Decidable ?P |- _ ] =>
    let Heqe := fresh "Heqe" in
    destruct (@Decidable_witness P H) eqn:Heqe
  end : decidability.

#[global]
Hint Extern 3 =>
  match goal with
    [ W : @Decidable_witness ?P ?H = true |- ?P ] =>
    exact (Decidable_sound P H W)
  end : decidability.

#[global]
Hint Extern 3 =>
  match goal with
    [ W : @Decidable_witness ?P ?H = false |- ~ ?P ] =>
    exact (Decidable_complete_alt P H W)
  end : decidability.

Lemma Decidable_witness_dec `{Decidable P} : P \/ ~ P.
Proof. decide P; firstorder. Qed.
Arguments Decidable_witness_dec {P _} /.

Lemma Decidable_witness_decides `{Decidable P} :
  if Decidable_witness then P else ~ P.
Proof. auto with decidability. Qed.
Arguments Decidable_witness_decides {P _} /.

Local Ltac t :=
  intros;
  repeat
    match goal with
      [ H : Decidable ?P |- _ ] =>
      let Heqe := fresh "Heqe" in
      destruct (@Decidable_witness P H) eqn:Heqe;
      [ apply (@Decidable_sound P H) in Heqe
      | apply (@Decidable_complete_alt P H) in Heqe ];
      clear H
    end;
  simpl; split; intros;
  try discriminate;
  try contradiction;
  try reflexivity;
  intuition.

Section DecidableLogic.

Local Obligation Tactic := t.

Global Program Instance not_Decidable {P : Prop} `{Decidable P} :
  Decidable (~ P) := {
  Decidable_witness := negb Decidable_witness
}.

Global Program Instance and_Decidable {P Q : Prop}
  `{Decidable P} `{Decidable Q} : Decidable (P /\ Q) := {
  Decidable_witness := andb (Decidable_witness (P:=P))
                            (Decidable_witness (P:=Q))
}.

Global Program Instance or_Decidable {P Q : Prop}
  `{Decidable P} `{Decidable Q} : Decidable (P \/ Q) := {
  Decidable_witness := orb (Decidable_witness (P:=P))
                           (Decidable_witness (P:=Q))
}.

Global Program Instance impl_Decidable {P Q : Prop}
  `{Decidable P} `{Decidable Q} : Decidable (P -> Q) := {
  Decidable_witness := if Decidable_witness (P:=P)
                       then Decidable_witness (P:=Q)
                       else true
}.

End DecidableLogic.

Local Ltac t' tac := t; apply tac; assumption.

Require Import Coq.Bool.Bool.

Global Program Instance bool_eq_Decidable {n m : bool} : Decidable (n = m) := {
  Decidable_witness := eqb n m
}.
Obligation 1. t' eqb_true_iff. Qed.

Require Import Coq.Strings.Ascii.
Require Import Coq.Bool.Sumbool.

Global Program Instance ascii_eq_Decidable {n m : Ascii.ascii} :
  Decidable (n = m) := {
  Decidable_witness := bool_of_sumbool (Ascii.ascii_dec n m)
}.
Obligation 1. t; destruct (Ascii.ascii_dec n m); auto; discriminate. Qed.

Require Import Coq.Arith.Arith.

Global Program Instance nat_eq_Decidable {n m : nat} : Decidable (n = m) := {
  Decidable_witness := beq_nat n m
}.
Obligation 1. t' beq_nat_true_iff. Qed.

Global Program Instance le_Decidable {n m} : Decidable (Nat.le n m) := {
  Decidable_witness := Compare_dec.leb n m
}.
Obligation 1. t' leb_iff. Qed.

Global Instance lt_Decidable {n m} : Decidable (lt n m) := le_Decidable.

Require Import Coq.NArith.NArith.

Global Program Instance N_eq_Decidable {n m : N} : Decidable (n = m) := {
  Decidable_witness := N.eqb n m
}.
Obligation 1. t' N.eqb_eq. Qed.

Global Program Instance Nle_Decidable {n m} : Decidable (N.le n m) := {
  Decidable_witness := N.leb n m
}.
Obligation 1. t' N.leb_le. Qed.

Global Program Instance Nlt_Decidable {n m} : Decidable (N.lt n m) := {
  Decidable_witness := N.ltb n m
}.
Obligation 1. t' N.ltb_lt. Qed.

Require Import Coq.ZArith.ZArith.

Global Program Instance Z_eq_Decidable {n m : Z} : Decidable (n = m) := {
  Decidable_witness := Z.eqb n m
}.
Obligation 1. t' Z.eqb_eq. Qed.

Global Program Instance Zle_Decidable {n m} : Decidable (Z.le n m) := {
  Decidable_witness := Z.leb n m
}.
Obligation 1. t' Z.leb_le. Qed.

Global Program Instance Zlt_Decidable {n m} : Decidable (Z.lt n m) := {
  Decidable_witness := Z.ltb n m
}.
Obligation 1. t' Z.ltb_lt. Qed.

Require Import Coq.QArith.QArith.

Global Program Instance Q_eq_Decidable {n m : Q} : Decidable (n == m) := {
  Decidable_witness := Qeq_bool n m
}.
Obligation 1. t' Qeq_bool_iff. Qed.

Global Program Instance Qle_Decidable {n m} : Decidable (Qle n m) := {
  Decidable_witness := Qle_bool n m
}.
Obligation 1. t' Qle_bool_iff. Qed.

Global Program Instance Qlt_Decidable {n m} : Decidable (Qlt n m) := {
  Decidable_witness := bool_of_sumbool (Qlt_le_dec n m)
}.
Obligation 1.
  t; destruct (Qlt_le_dec n m); simpl in H; auto.
    discriminate.
  apply Qle_not_lt in H; auto.
Qed.

Section ListDec.
  Variable A : Type.
  Variable eq_dec : forall x y : A, {x = y} + {x <> y}.

  Global Program Instance list_eq_Decidable {xs ys : list A} :
    Decidable (xs = ys) := {
    Decidable_witness := bool_of_sumbool (List.list_eq_dec eq_dec xs ys)
  }.
  Obligation 1.
    t; destruct (List.list_eq_dec eq_dec xs ys);
    simpl in H; auto; discriminate.
  Qed.

  Global Program Instance list_in_Decidable {xs : list A} {x : A} :
    Decidable (List.In x xs) := {
    Decidable_witness := List.existsb (fun y => bool_of_sumbool (eq_dec x y)) xs
  }.
  Obligation 1.
    remember (fun _ => proj1_sig _) as F; t;
    destruct (List.existsb_exists F xs);
    destruct (List.existsb F xs);
    auto; try discriminate.
      destruct (H0 H).
      destruct H2.
      rewrite HeqF in H3.
      destruct (eq_dec x x0);
      subst; auto; discriminate.
    apply H1.
    exists x; split; auto.
    rewrite HeqF.
    destruct (eq_dec x x); auto.
  Qed.
End ListDec.


Inductive color := Green | Yellow | Red.

Definition min_color a b :=
  match a, b with
  | Red, _ | _, Red => Red
  | Yellow, _ | _, Yellow => Yellow
  | _, _ => Green
  end.


From Hammer Require Import Hammer.

Set Implicit Arguments.

Require Import Le.
Require Import Lt.
Require Import Arith.
Require Import Recdef.
Require Coq.Program.Wf.
Require List.

Section measure_wf.

Variables T M: Set.
Variable R: M -> M -> Prop.
Hypothesis wf: well_founded R.
Variable m: T -> M.

Definition MR (x y: T): Prop := R (m x) (m y).

Lemma measure_wf: well_founded MR.
Proof with auto. hammer_hook "fix_measure_utils" "fix_measure_utils.measure_wf".
unfold well_founded.
cut (forall a: M, forall a0: T, m a0 = a -> Acc MR a0).
intros.
apply (H (m a))...
apply (@well_founded_ind M R wf (fun mm => forall a, m a = mm -> Acc MR a)).
intros.
apply Acc_intro.
intros.
unfold MR in H1.
rewrite H0 in H1.
apply (H (m y))...
Defined.

End measure_wf.

Section Fix_measure_rects.

Variable A: Set.
Variable R: A -> A -> Prop.
Variable Rwf: well_founded R.
Variable P: A -> Type.
Variable f: forall x: A, (forall y: { y: A | R y x }, P (proj1_sig y)) -> P x.

Lemma F_unfold x r:
Wf.Fix_F_sub A R P f x r =
f (fun y => Wf.Fix_F_sub A R P f (proj1_sig y) (Acc_inv r (proj2_sig y))).
Proof. hammer_hook "fix_measure_utils" "fix_measure_utils.F_unfold". intros. case r; auto. Qed.

Lemma F_sub_rect
(Q: forall x, P x -> Type)
(inv: forall x: A,
(forall (y: A) (H: R y x) (a: Acc R y),
Q y (Wf.Fix_F_sub A R P f y a)) ->
forall (a: Acc R x),
Q x (f (fun y: {y: A | R y x} =>
Wf.Fix_F_sub A R P f (proj1_sig y) (Acc_inv a (proj2_sig y)))))
: forall x a, Q _ (Wf.Fix_F_sub A R P f x a).
Proof with auto. hammer_hook "fix_measure_utils" "fix_measure_utils.F_sub_rect".
set (R' := fun (x: A) => forall a, Q _ (Wf.Fix_F_sub A R P f x a)).
cut (forall x, R' x)...
apply (well_founded_induction_type Rwf).
subst R'.
simpl.
intros.
rewrite F_unfold...
Qed.

Hypothesis equiv_lowers:
forall x0 (g h: forall x: {y: A | R y x0}, P (proj1_sig x)),
(forall x p p', g (exist (fun y: A => R y x0) x p) = h (exist _ x p')) ->
f g = f h.

Lemma eq_F_sub x: forall (a a': Acc R x),
Wf.Fix_F_sub A R P f x a =
Wf.Fix_F_sub A R P f x a'.
Proof. hammer_hook "fix_measure_utils" "fix_measure_utils.eq_F_sub".
intro a.
pattern x, (Wf.Fix_F_sub A R P f x a).
apply F_sub_rect.
intros.
rewrite F_unfold.
apply equiv_lowers.
intros.
apply H.
assumption.
Qed.

Lemma unfold x:
Wf.Fix_sub A R Rwf P f x =
f (fun y => Wf.Fix_sub A R Rwf P f (proj1_sig y)).
Proof. hammer_hook "fix_measure_utils" "fix_measure_utils.unfold". intros. unfold Wf.Fix_sub.
rewrite F_unfold.
apply equiv_lowers.
intros.
apply eq_F_sub.
Qed.

Lemma rect
(Q: forall x, P x -> Type)
(inv: forall
(x: A)
(H: forall (y: A), R y x -> Q y (Wf.Fix_sub A R Rwf P f y))
(a: Acc R x),
Q x (f (fun y: {y: A | R y x} => Wf.Fix_sub A R Rwf P f (proj1_sig y))))
: forall x, Q _ (Wf.Fix_sub A R Rwf P f x).
Proof with auto. hammer_hook "fix_measure_utils" "fix_measure_utils.rect".
unfold Wf.Fix_sub.
intros.
apply F_sub_rect.
intros.
assert (forall y: A, R y x0 -> Q y (Wf.Fix_F_sub A R P f y (Rwf y)))...
set (inv x0 X0 a). clearbody q.
rewrite <- (equiv_lowers (fun y: {y: A | R y x0} => Wf.Fix_F_sub A R P f (proj1_sig y) (Rwf (proj1_sig y))) (fun y: {y: A | R y x0} => Wf.Fix_F_sub A R P f (proj1_sig y) (Acc_inv a (proj2_sig y))))...
intros.
apply eq_F_sub.
Qed.

End Fix_measure_rects.

Module Example.

Import List.

Definition tail (l: list nat): list nat :=
match l with
| nil => nil
| _ :: t => t
end.

Program Fixpoint bla (l: list nat) {measure (length l) lt}: nat :=
match l with
| nil => 3
| _ => bla (tail l) + 2
end.

Next Obligation.
destruct l.
elimtype False. apply H. reflexivity.
simpl. apply le_refl.
Qed.

Goal forall x, 3 <= bla x.
Proof with auto. hammer_hook "fix_measure_utils" "fix_measure_utils.Example.tail".
intro x.
pattern (bla x).
set (fun n : nat => 3 <= n).
unfold bla.
generalize x. clear x.
apply rect; intros.
destruct x0...
subst P.
simpl.
destruct x...
replace 3 with (3 + 0)...
apply plus_le_compat...
apply H.
unfold Wf.MR...
Qed.

End Example.


(** * Additional theorems and definitions *)

Require Import UniMath.Foundations.Preamble.
Require Import UniMath.MoreFoundations.Tactics.

Require Export UniMath.Topology.Prelim.

(** ** for RationalNumbers.v *)

Require Export UniMath.NumberSystems.RationalNumbers.
Require Export UniMath.Algebra.Archimedean.

Local Open Scope hq_scope.

Notation "x <= y" := (hqleh x y) : hq_scope.
Notation "x >= y" := (hqgeh x y) : hq_scope.
Notation "x < y" := (hqlth x y) : hq_scope.
Notation "x > y" := (hqgth x y) : hq_scope.
Notation "/ x" := (hqmultinv x) : hq_scope.
Notation "x / y" := (hqdiv x y) : hq_scope.
Notation "2" := (hztohq (nattohz 2)) : hq_scope.

Lemma hzone_neg_hzzero : neg (1%hz = 0%hz).
Proof.
  confirm_not_equal isdeceqhz.
Qed.

Definition one_intdomnonzerosubmonoid : intdomnonzerosubmonoid hzintdom.
Proof.
  exists 1%hz ; simpl.
  exact hzone_neg_hzzero.
Defined.

Opaque hz.

Lemma hq2eq1plus1 : 2 = 1 + 1.
Proof.
  confirm_equal isdeceqhq.
Qed.

Lemma hq2_gt0 : 2 > 0.
Proof.
  confirm_yes hqgthdec 2 0.
Qed.

Lemma hq1_gt0 : 1 > 0.
Proof.
  confirm_yes hqgthdec 1 0.
Qed.

Lemma hq1ge0 : (0 <= 1)%hq.
Proof.
  confirm_yes hqlehdec 0 1.
Qed.

Lemma hqgth_hqneq :
  ∏ x y : hq, x > y -> hqneq x y.
Proof.
  intros x y Hlt Heq.
  rewrite Heq in Hlt.
  now apply isirreflhqgth with y.
Qed.

Lemma hqldistr :
  ∏ x y z, x * (y + z) = x * y + x * z.
Proof.
  intros x y z.
  now apply ringldistr.
Qed.

Lemma hqmult2r :
  ∏ x : hq, x * 2 = x + x.
Proof.
  intros x.
  now rewrite hq2eq1plus1, hqldistr, (hqmultr1 x).
Qed.

Lemma hqplusdiv2 : ∏ x : hq, x = (x + x) / 2.
  intros x.
  apply hqmultrcan with 2.
  - now apply hqgth_hqneq, hq2_gt0.
  - unfold hqdiv.
    rewrite hqmultassoc.
    rewrite (hqislinvmultinv 2).
    2: now apply hqgth_hqneq, hq2_gt0.
    rewrite (hqmultr1 (x + x)).
    apply hqmult2r.
Qed.

Lemma hqlth_between :
  ∏ x y : hq, x < y -> total2 (λ z, (x < z) × (z < y)).
Proof.
  assert (H0 : / 2 > 0).
  { apply hqgthandmultlinv with 2.
    - apply hq2_gt0.
    - rewrite hqisrinvmultinv, hqmultx0.
      + now apply hq1_gt0.
      + now apply hqgth_hqneq, hq2_gt0.
  }
  intros x y Hlt.
  exists ((x + y) / 2).
  split.
  - pattern x at 1.
    rewrite (hqplusdiv2 x).
    unfold hqdiv.
    apply (hqlthandmultr _ _ (/ 2)).
    + exact H0.
    + now apply (hqlthandplusl _ _ x Hlt).
  - pattern y at 2.
    rewrite (hqplusdiv2 y).
    unfold hqdiv.
    apply (hqlthandmultr _ _ (/ 2)).
    + exact H0.
    + now apply (hqlthandplusr _ _ y Hlt).
Qed.

Lemma hq0lehandplus:
  ∏ n m : hq,
    0 <= n -> 0 <= m -> 0 <= (n + m).
Proof.
  intros n m Hn Hm.
  eapply istranshqleh, hqlehandplusl, Hm.
  now rewrite hqplusr0.
Qed.

Lemma hq0lehandmult:
  ∏ n m : hq, 0 <= n -> 0 <= m -> 0 <= n * m.
Proof.
  intros n m.
  exact hqmultgeh0geh0.
Qed.

Lemma hq0leminus :
  ∏ r q : hq, r <= q -> 0 <= q - r.
Proof.
  intros r q Hr.
  apply hqlehandplusrinv with r.
  unfold hqminus.
  rewrite hqplusassoc, hqlminus.
  now rewrite hqplusl0, hqplusr0.
Qed.

Lemma hqinv_gt0 (x : hq) : 0 < x → 0 < / x.
Proof.
  unfold hqlth.
  intros Hx.
  apply hqgthandmultlinv with x.
  - exact Hx.
  - rewrite hqmultx0.
    rewrite hqisrinvmultinv.
    + exact hq1_gt0.
    + apply hqgth_hqneq.
      exact Hx.
Qed.

Lemma hztohqandleh':
  ∏ n m : hz, (hztohq n <= hztohq m)%hq → hzleh n m.
Proof.
  intros n m Hle Hlt.
  simple refine (Hle _).
  apply hztohqandgth.
  exact Hlt.
Qed.
Lemma hztohqandlth':
  ∏ n m : hz, (hztohq n < hztohq m)%hq -> hzlth n m.
Proof.
  intros n m Hlt.
  apply neghzgehtolth.
  intro Hle.
  apply hztohqandgeh in Hle.
  apply hqgehtoneghqlth in Hle.
  apply Hle.
  exact Hlt.
Qed.

Close Scope hq_scope.