(* -------------------------------------------------------------------- *)
require import AllCore List StdBigop StdOrder IntDiv Distr Finite.
(*---*) import IntOrder Bigint MUniform Range.

(* -------------------------------------------------------------------- *)
(* We keep intervals closed for now, to justify attaching the [_.._]
   notation to this distribution. *)
op dinter (i j : int) = duniform (range i (j + 1)).

lemma dinter1E (i j : int) x:
  mu1 (dinter i j) x = if (i <= x <= j) then 1%r/(j - i + 1)%r else 0%r.
proof. by rewrite duniform1E mem_range undup_id 1:range_uniq size_range /#. qed.

lemma dinterE (i j : int) P:
  mu (dinter i j) P
  = (size (filter P (range i (j + 1))))%r / (max 0 (j - i + 1))%r.
proof. by rewrite duniformE undup_id 1:range_uniq size_range size_filter /#. qed.

lemma weight_dinter (i j : int):
  weight (dinter i j) = b2r (i <= j).
proof. by rewrite dinterE filter_predT size_range /#. qed.

lemma supp_dinter (i j : int) x:
  x \in (dinter i j) <=> i <= x <= j.
proof. by rewrite /support dinter1E; case (i <= x <= j)=> //= /#. qed.

lemma supp_dinter1E (x : int) (i j : int) :
  x \in (dinter i j) => mu1 (dinter i j) x = 1%r / (j - i + 1)%r.
proof. by rewrite supp_dinter dinter1E => ->. qed.

lemma dinter_ll (i j : int): i <= j => is_lossless (dinter i j).
proof. move=> Hij;apply /drange_ll => /#. qed.

lemma dinter_uni (i j : int): is_uniform (dinter i j).
proof. apply drange_uni. qed.

lemma finite_dinter (i j : int) : is_finite (support (dinter i j)).
proof.
rewrite is_finiteE; exists (range i (j+1)).
by rewrite range_uniq /= => x; rewrite mem_range supp_dinter /#.
qed.

lemma perm_eq_dinter (i j : int) : 
  perm_eq (to_seq (support (dinter i j))) (range i (j+1)).
proof. 
apply: uniq_perm_eq; first exact/uniq_to_seq/finite_dinter.
- exact: range_uniq.
by move=> x; rewrite mem_to_seq ?finite_dinter // supp_dinter mem_range /#.
qed.

lemma perm_eq_dinter_pred (i j : int) : 
    perm_eq (to_seq (support (dinter i (j-1)))) (range i j).
proof. by have /# := perm_eq_dinter i (j-1). qed.  

(* -------------------------------------------------------------------- *)
lemma duni_range_dvd (p q : int) : 0 < p => 0 < q => q %| p =>
  dmap (dinter 0 (p-1)) (fun x => x %% q) = dinter 0 (q-1).
proof.
move=> gt0_p gt0_q dvd_qp; apply/eq_distr=> i; apply/eq_sym.
rewrite dinter1E ler_subr_addl (addrC 1) -ltzE -mem_range.
case: (i \in range 0 q) => /= [/mem_range [ge0i lei]|Nrgi]; last first.
- pose d := dmap _ _; suff /supportPn/eq_sym //: !(i \in d).
  apply: contra Nrgi => /supp_dmap[j [/supp_dinter [ge0j lej] /=]].
  by move=> ->>; rewrite mem_range modz_ge0 1:gtr_eqF //= ltz_pmod.
pose s := flatten (map
  (fun k => map (fun i : int => k * q + i) (range 0 q)) (range 0 (p %/ q))).
have uniq_s: uniq s; first apply: uniq_flatten_map.
- move=> j /=; rewrite map_inj_in_uniq ?range_uniq.
  by move => x y _ _ /=; apply/addrI.
- move=> x y /mem_range rg_x /mem_range rg_y /hasP /=.
  case=> j [/mapP[/= k [/mem_range rg_k]] ->>].
  by case/mapP=> l [/mem_range rg_l] {rg_x} {rg_y} /#.
- by apply: range_uniq.
have mem_s: forall x, (x \in s) <=> (x \in range 0 p).
- move=> x; rewrite mem_range; split.
  - case/flatten_mapP=> j [/= /mem_range rgj] /=.
    case/mapP => [k [/mem_range rgk]] ->>; smt(@IntDiv).
  - case=> ge0x lex; apply/flatten_mapP => /=.
    exists (x %/ q); rewrite mem_range.
    rewrite divz_ge0 // ge0x /=; split; 1: smt(@IntDiv).
    apply/mapP; exists (x %% q); rewrite mem_range.
    rewrite modz_ge0 1:gtr_eqF //= ltz_pmod //= &(divz_eq).
have eqs: perm_eq s (range 0 p).
- by apply: uniq_perm_eq => //; apply/range_uniq.
rewrite dmap1E duniformE /= undup_id ?range_uniq /pred1 /(\o).
rewrite size_range /max gt0_p /=; have := perm_eqP s (range 0 p).
case=> + _ - /(_ eqs) <- @/s; rewrite count_flatten sumzE /=.
pose t := map _ _; rewrite BIA.big_seq -(BIA.eq_bigr _ (fun _ => 1)) /=.
- move=> k /mapP[xs [/mapP[j [/mem_range [ge0j lej]] ->>]] ->>].
  rewrite count_map /preim -(eq_in_count (pred1 i)) => [k|].
  - by case/mem_range=> [ge0k lek] /=; rewrite modzMDl pmod_small.
  by rewrite eq_sym count_uniq_mem ?range_uniq mem_range ge0i lei.
rewrite -(BIA.big_seq _ t) big_constz count_predT !size_map.
rewrite size_range /= /max; case/dvdzP: dvd_qp => r ->>.
have gt0_r : 0 < r by apply: (pmulr_lgt0 q).
rewrite mulzK 1:gtr_eqF // fromintM gt0_r /=.
rewrite  RField.invrM ?eq_fromint 1,2:gtr_eqF //.
by rewrite RField.mulrCA RField.divff // eq_fromint gtr_eqF.
qed.

(* -------------------------------------------------------------------- *)
abstract theory Cost.
  op cdinterval : int -> int.
  axiom ge0_cdinterval m : 0 <= cdinterval m.
  
  schema cost_dinterval {i j : int} (k:int) : 
    cost [ i <= j <= k - i : dinter i (j - 1)] = 
    cost [true : i] + cost [true : j] + N (cdinterval k).
  hint simplify cost_dinterval.
end Cost.