require import AllCore Discrete Finite Distr SmtMap List.

op duni ['a] = duniform (Finite.to_seq<:'a> predT).

lemma duni_full : is_finite predT<:'a> => is_full duni<:'a>.
proof. by move => finT x @/duni; rewrite supp_duniform mem_to_seq. qed.

lemma duni_uni : is_uniform duni<: 'a> by exact duniform_uni.

lemma duni_ll : is_finite predT<:'a> => is_lossless duni<: 'a>.
proof.
move => finT; apply duniform_ll.
suff : witness \in to_seq predT<:'a>; smt(mem_to_seq).
qed.

lemma supp_duni : is_finite (predT<:'a>) => support duni = predT<:'a>.
proof. smt(duni_full). qed.

lemma finite_duni : is_finite (support duni<:'a>) by exact finite_duniform.

lemma finite_for_size (p : 'a -> bool) s :
  is_finite_for p s => size (to_seq p) = size s.
proof.
case => s_uniq mem_s; have ? : is_finite p by exists s.
apply/perm_eq_size/uniq_perm_eq; 2,3: smt(mem_to_seq).
apply/uniq_to_seq; by exists s.
qed.

lemma size_dprod (d1 : 'a distr) (d2 : 'b distr) :
  is_finite (support d1) => is_finite (support d2) =>
  size (to_seq (support (d1 `*` d2))) =
  size (to_seq (support d1)) * size (to_seq (support d2)).
proof.
move => [s1 fin_for_s1] [s2 fin_for_s2].
rewrite (finite_for_size _ _ fin_for_s1) (finite_for_size _ _ fin_for_s2).
suff fin_for_s :
  is_finite_for (support (d1 `*` d2)) (allpairs (fun x y => (x,y)) s1 s2).
- by rewrite (finite_for_size _ _ fin_for_s) size_allpairs.
split => [|x]; first by apply allpairs_uniq => /#.
rewrite allpairsP; smt(supp_dprod).
qed.

(* types whose cardinality is a power of two (i.e. types that can be
represented by fixing-length bitstrings *)
op fixed_type ['a] = 
  finite_type<:'a> /\ exists n, 0 <= n /\ size (to_seq predT<:'a>) = 2^n.

lemma fixed_fin ['a] : fixed_type<:'a> => finite_type<:'a> by auto => />.

lemma prod_fixed ['a 'b] : 
  fixed_type<:'a> => fixed_type<:'b> => fixed_type<:'a * 'b>.
proof.
move => [fin_a [na [na0 pow_a]]] [fin_b [nb [nb0 pow_b]]].
have ab_fin : finite_type<:'a*'b> by apply finite_type_pair.
split => //; move: fin_a fin_b => [sa ffa] [sb ffb].
exists (na+nb).
suff -> : size (to_seq predT<:'a*'b>) = 
         size (allpairs (fun (x : 'a) (y : 'b) => (x, y)) sa sb).
- rewrite size_allpairs -(finite_for_size _ _ ffa) -(finite_for_size _ _ ffb).
  smt(exprD_nneg size_ge0).
apply uniq_size_uniq => //. 
+ by apply allpairs_uniq; smt(). 
+ move => [x y]. rewrite allpairsP mem_to_seq // /predT /=. 
  exists (x,y); smt().
+ by apply uniq_to_seq.
qed.

(* This is the minimum we need to know to encode tuples.
   We may want to substitute this with something more concrete *)
abstract theory Bitstring.

type bitstring.

axiom bitstring_inf : ! finite_type<:bitstring>.
axiom bitstring_count : countableT<:bitstring>.

op encode : 'a -> bitstring.
op decode : bitstring -> 'a.

axiom encodeK ['a] : countableT<:'a> => cancel encode<:'a> decode.

end Bitstring.

pred commutative ['a 'b] (f : 'a * 'a -> 'b) =
  forall (x y : 'a), f(x, y) = f(y, x).
pred associative ['a] (f : 'a * 'a -> 'a) =
  forall (x y z : 'a), f(x, f(y, z)) = f(f(x, y), z).
pred neutral ['a] (f : 'a * 'a -> 'a) (n : 'a) =
  forall (x : 'a), f(x, n) = x /\ f(n, x) = x.
pred group ['a] (f : 'a * 'a -> 'a) (inv : 'a -> 'a) (e : 'a) =
     associative f
  /\ neutral f e
  /\ (forall x, f(x, inv x) = e)
  /\ (forall x, f(inv x, x) = e).
pred cancel ['a] (f : 'a * 'a -> 'a) (n : 'a) =
  forall (x : 'a), f(x, x) = n.