(* -------------------------------------------------------------------- *)
require import AllCore List Distr StdOrder.
require import FinType.
require (*--*) Tuple.
(*---*) import IntOrder RealOrder.

(* -------------------------------------------------------------------- *)
clone import FinType as Alphabet.

op n : {int | 0 <= n} as ge0_n.

type word.

op mkword : t list -> word.
op ofword : word -> t list.
op dchar  : t.

op wordw : t list = mkseq (fun _ => dchar) n.

(* -------------------------------------------------------------------- *)
lemma nosmt size_wordw: size wordw = n.
proof. by rewrite size_mkseq ler_maxr ?ge0_n. qed.

axiom nosmt mkwordK : cancel ofword mkword.

axiom nosmt ofwordK :
  forall (s : t list), size s = n =>
    ofword (mkword s) = s.

axiom nosmt mkword_out :
  forall (s : t list), size s <> n =>
    ofword (mkword s) = wordw.

lemma nosmt size_word : forall w, size (ofword w) = n.
proof.
move=> w; rewrite -mkwordK; case: (size (ofword w) = n).
by rewrite mkwordK. by move/mkword_out=> ->; rewrite size_wordw.
qed.

lemma nosmt mkwordW (P : word -> bool):
     (forall (s : t list), size s = n => P (mkword s))
  => forall n, P n.
proof. by move=> ih n; rewrite -mkwordK; apply/ih/size_word. qed.

lemma nosmt ofword_inj : injective ofword.
proof. by apply/(can_inj _ _ mkwordK). qed.

lemma nosmt mkword_pinj (s1 s2 : t list) :
  size s1 = n => size s2 = n =>
  mkword s1 = mkword s2 =>
  s1 = s2.
proof.
move=> size_s1 size_s2 eq_mkword.
by rewrite -(@ofwordK s1) // -(@ofwordK s2) // eq_mkword.
qed.

(* -------------------------------------------------------------------- *)
op "_.[_]" (w : word) (i : int): t =
  if 0 <= i < n then nth witness (ofword w) i else dchar
axiomatized by getE.

op "_.[_<-_]" (w : word) (i : int) (x : t) : word =
  mkword (mkseq (fun k => if i = k then x else w.[k]) n)
axiomatized by setE.

(* -------------------------------------------------------------------- *)
lemma wordP (w1 w2 : word):
  w1 = w2 <=> forall i, 0 <= i < n => w1.[i] = w2.[i].
proof.
split=> [->|] //= h; apply/(inj_eq _ ofword_inj)/(eq_from_nth witness).
+ by rewrite !size_word.
by move=> i; rewrite size_word=> ^i_bound /h; rewrite !getE i_bound.
qed.

(* -------------------------------------------------------------------- *)
lemma wordW (P : word -> bool):
     (forall (s : t list), size s = n => P (mkword s))
  => forall w, P w.
proof. by move=> ih n; rewrite -mkwordK; apply/ih/size_word. qed.

(* -------------------------------------------------------------------- *)
lemma get_set_if (w : word) (x : t) (i j : int):
  w.[i <- x].[j] = if 0 <= i < n /\ j = i then x else w.[j].
proof.
rewrite !Self.getE !Self.setE ofwordK ?size_mkseq ?ler_maxr ?ge0_n //.
case: (0 <= j < n)=> j_n.
+ rewrite nth_mkseq ?j_n /= (eq_sym i) getE.
  by case: (j = i)=> [<-|]; rewrite j_n.
by case: (j = i)=> [<-|]; rewrite ?j_n.
qed.

(* -------------------------------------------------------------------- *)
lemma get_set (w : word) (x : t) (i j : int): 0 <= i < n =>
  w.[i <- x].[j] = if j = i then x else w.[j].
proof. by move=> lt_in; rewrite get_set_if lt_in. qed.

(* -------------------------------------------------------------------- *)
lemma set_out (i : int) (x : t) (w : word):
  ! (0 <= i < n) => w.[i <- x] = w.
proof.
by move=> Nlt_in; apply/wordP=> j _; rewrite get_set_if Nlt_in.
qed.

(* -------------------------------------------------------------------- *)
lemma set_neg (i : int) (a : t) (w : word):
  i < 0 => w.[i<- a] = w.
proof. by move=> lt0_i; rewrite set_out // lezNgt lt0_i. qed.

(* -------------------------------------------------------------------- *)
lemma set_above (i : int) (a : t) (w : word):
  n <= i => w.[i <- a] = w.
proof. by move=> le_ni; rewrite set_out // ltzNge le_ni. qed.

(* -------------------------------------------------------------------- *)
lemma set_set_if (w : word) (k k' : int) (x x' : t):
       w.[k <- x].[k' <- x']
    =  if   k = k'
       then w.[k' <- x']
       else w.[k' <- x'].[k <- x].
proof.
apply/wordP=> i i_n; rewrite !Self.setE !Self.getE i_n /=.
rewrite (fun_if ofword) (fun_if (nth witness)) if_arg.
rewrite !ofwordK ?size_mkseq ?ler_maxr ?ge0_n //.
rewrite !nth_mkseq //= !getE !ofwordK ?size_mkseq ?ler_maxr ?ge0_n // i_n /=.
case: (k = k')=> [->|].
+ by case: (k' = i)=> //= k'_neq_i; rewrite nth_mkseq //= getE i_n k'_neq_i.
by rewrite !nth_mkseq //=; case: (k' = i); case: (k = i).
qed.

(* -------------------------------------------------------------------- *)
lemma set_set_eq (w : word) (k : int) (x x' : t):
  w.[k <- x].[k <- x'] = w.[k <- x'].
proof. by rewrite set_set_if. qed.

(* -------------------------------------------------------------------- *)
lemma set_set_swap (w : word) (k k' : int) (x x' : t):
  k <> k => w.[k <- x].[k' <- x'] = w.[k' <- x'].[k' <- x'].
proof. by rewrite set_set_if. qed.

(* -------------------------------------------------------------------- *)
op offun f: word = mkword (mkseq f n).

lemma offunifE f i:
  (offun f).[i] = if 0 <= i < n then f i else dchar.
proof.
rewrite getE ofwordK ?size_mkseq ?size_word ?ler_maxr ?ge0_n//.
by case: (0 <= i < n)=> //= h; rewrite nth_mkseq.
qed.

lemma offunE f i: 0 <= i < n => (offun f).[i] = f i.
proof. by move=> lt_in; rewrite offunifE lt_in. qed.

lemma ofwordKmkseq (f : int -> t):
  ofword (offun f) = mkseq f n.
proof.
by rewrite ofwordK 1:size_mkseq 1:StdOrder.IntOrder.ler_maxr 1:ge0_n.
qed.

(* -------------------------------------------------------------------- *)
op map (f : t -> t) (w : word) : word = mkword (map f (ofword w)).

lemma mapE f w i: 0 <= i < n => (map f w).[i] = f w.[i].
proof.
rewrite !Self.getE ofwordK ?size_map 1:size_word// => ^h -> /=.
by apply/nth_map; rewrite size_word.
qed.

(* -------------------------------------------------------------------- *)
clone Tuple as Enum with
  type t <- t,
    op Support.enum <- enum
proof * by exact/enum_spec.

(* -------------------------------------------------------------------- *)
clone include FinType with
  type t    <- word,
  op enum = List.map mkword (Enum.wordn n)
proof * rename [op] "enum" as "words".

realize enum_spec.
proof.
move=> x; rewrite count_map -(Enum.word_spec n (ofword x) _).
+ by rewrite size_word ler_maxr 1:ge0_n.
apply/eq_in_count=> s /Enum.wordnP; rewrite !pred1E /preim ler_maxr 1:ge0_n.
by move=> sx; rewrite -(can_eq _ _ mkwordK) ofwordK.
qed.

(* -------------------------------------------------------------------- *)
lemma word_card: card = Alphabet.card^n.
proof. by rewrite /card size_map Enum.size_wordn ler_maxr ?ge0_n. qed.

(* -------------------------------------------------------------------- *)
clone MFinite as DWord with
  type t <- word,
  op   Support.enum = List.map mkword (Enum.wordn n),
  op   Support.card <- card
proof Support.enum_spec by exact/enum_spec.

(* -------------------------------------------------------------------- *)
theory ListSample.
require import DList.

clone MFinite as DAlphabet with
  type         t <= Alphabet.t,
  theory Support <= Alphabet
proof *.

lemma dword_is_dlist:
  dmap (dlist DAlphabet.dunifin n) mkword = DWord.dunifin.
proof.
apply: eq_funi_ll.
+ apply: is_full_funiform.
  + apply: dmap_fu_in=> w; exists (ofword w).
    rewrite mkwordK //= (supp_dlist _ _ _ ge0_n) size_word //=.
    by apply: allP=> l _; exact: DAlphabet.supp_dunifin.
  apply: dmap_uni_in_inj.
  + move=> ls ls'; rewrite !(supp_dlist _ _ _ ge0_n)=> /> + _ + _.
    by move=> /ofwordK {2}<- /ofwordK {2}<- ->.
  exact/dlist_uni/DAlphabet.dunifin_uni.
+ exact/dmap_ll/dlist_ll/DAlphabet.dunifin_ll.
+ exact/DWord.dunifin_funi.
exact/DWord.dunifin_ll.
qed.

end ListSample.