(* --------------------------------------------------------------------
 * Copyright (c) - 2012-2015 - IMDEA Software Institute and INRIA
 * Distributed under the terms of the CeCILL-C license
 * -------------------------------------------------------------------- *)

(* This API has been mostly inspired from the [bigop] library of the
 * ssreflect Coq extension. *)

pragma +implicits.

(* -------------------------------------------------------------------- *)
require import Fun Pred Pair Int IntExtra NewList Ring.
require (*--*) NewMonoid.

import Ring.IntID.

(* -------------------------------------------------------------------- *)
type t.

clone export NewMonoid as Support with type t <- t.

(* -------------------------------------------------------------------- *)
op big (P : 'a -> bool) (F : 'a -> t) (r : 'a list) =
  foldr Support.(+) idm (map F (filter P r)).

(* -------------------------------------------------------------------- *)
op bigi (P : int -> bool) (F : int -> t) i j =
  big P F (range i j).

(* -------------------------------------------------------------------- *)
lemma big_nil (P : 'a -> bool) (F : 'a -> t): big P F [] = idm.
proof. by []. qed.

(* -------------------------------------------------------------------- *)
lemma big_cons (P : 'a -> bool) (F : 'a -> t) x s:
  big P F (x :: s) = if P x then F x + big P F s else big P F s.
proof. by rewrite {1}/big /= @(fun_if (map F)); case (P x). qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_rec (K : t -> bool) r P (F : 'a -> t):
  K idm => (forall i x, P i => K x => K (F i + x)) => K (big P F r).
proof.
  move=> K0 Kop; elim: r => //= i r; rewrite big_cons.
  by case (P i) => //=; apply/Kop.
qed.

lemma nosmt big_ind (K : t -> bool) r P (F : 'a -> t):
     (forall x y, K x => K y => K (x + y))
  => K idm => (forall i, P i => K (F i))
  => K (big P F r).
proof.
  move=> Kop Kidx K_F; apply/big_rec => //.
  by move=> i x Pi Kx; apply/Kop => //; apply/K_F.
qed.

lemma nosmt big_rec2:
  forall (K : t -> t -> bool) r P (F1 F2 : 'a -> t),
       K idm idm
    => (forall i y1 y2, P i => K y1 y2 => K (F1 i + y1) (F2 i + y2))
    => K (big P F1 r) (big P F2 r).
proof.
  move=> K r P F1 F2 KI KF; elim: r => //= i r IHr.
  by rewrite !big_cons; case (P i) => _ //=; apply/KF.
qed.

lemma nosmt big_ind2:
  forall (K : t -> t -> bool) r P (F1 F2 : 'a -> t),
       (forall x1 x2 y1 y2, K x1 x2 => K y1 y2 => K (x1 + y1) (x2 + y2))
    => K idm idm
    => (forall i, P i => K (F1 i) (F2 i))
    => K (big P F1 r) (big P F2 r).
proof.
  move=> K r P F1 F2 Kop KI KF; apply/big_rec2 => //.
  by move=> i x1 x2 Pi Kx1x2; apply/Kop => //; apply/KF.
qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_endo (f : t -> t):
     f idm = idm
  => (forall (x y : t), f (x + y) = f x + f y)
  => forall r P (F : 'a -> t),
       f (big P F r) = big P (f \o F) r.
proof.
  (* FIXME: should be a consequence of big_morph *)
  move=> fI fM; elim=> //= i r IHr P F; rewrite !big_cons.
  by case (P i) => //=; rewrite 1?fM IHr.
qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_map ['a 'b] (h : 'b -> 'a) (P : 'a -> bool) F s:
  big P F (map h s) = big (P \o h) (F \o h) s.
proof. by elim: s => // x s; rewrite map_cons !big_cons=> ->. qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_comp ['a] (h : t -> t) (P : 'a -> bool) F s:
  h idm = idm => morphism_2 h Support.(+) Support.(+) =>
  h (big P F s) = big P (h \o F) s.
proof. 
  move=> Hidm Hh;elim: s => // x s; rewrite !big_cons => <-.
  by rewrite /(\o) -Hh;case (P x) => //.
qed.

lemma nosmt big_Dl ['a] (h : t -> t -> t) (P : 'a -> bool) F s t:
   left_zero idm h => 
   left_distributive h Support.(+) =>
   h (big P F s) t = big P (fun a => h (F a) t) s.
proof.
  move=> Hlz HDl.
  cut /= -> := big_comp (fun t1 => h t1 t) P F s _ _=> //=.
  + by apply Hlz.
  move=> t1 t2; apply (HDl t1 t2 t).
qed.

lemma nosmt big_Dr ['a] (h:t -> t -> t) (P:'a -> bool) F s t:
   right_zero idm h => 
   right_distributive h Support.(+) =>
   h t (big P F s) = big P (fun a => h t (F a)) s.
proof.
  move=> Hrz HDr.
  by cut /= -> := big_comp (h t) P F s _ _=> //=;[apply Hrz|apply (HDr t)]. 
qed.

lemma nosmt big_D ['a 'b] (h : t -> t -> t) (P1 : 'a -> bool) F1 s1
                                             (P2 : 'b -> bool) F2 s2:
   commutative h =>
   left_zero idm h => 
   left_distributive h Support.(+) =>
   h (big P1 F1 s1) (big P2 F2 s2) =
     big P1 (fun a1 => big P2 (fun a2 => h (F1 a1) (F2 a2)) s2) s1.
proof.
  move=> Hc Hlz HDl.
  rewrite big_Dl => //; congr;apply fun_ext=> a1.
  rewrite big_Dr //;smt.
qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_nth x0 (P : 'a -> bool) (F : 'a -> t) s:
  big P F s = big (P \o (nth x0 s)) (F \o (nth x0 s)) (iota_ 0 (size s)).
proof. by rewrite -{1}@(mkseq_nth x0 s) /mkseq big_map. qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_const (P : 'a -> bool) x s:
  big P (fun i => x) s = iter (count P s) ((+) x) idm.
proof.
  elim: s=> [|y s ih]; [by rewrite iter0 | rewrite big_cons /=].
  by rewrite ih; case (P y) => //; rewrite addzC iterS // count_ge0.
qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_seq1 (F : 'a -> t) x: big predT F [x] = F x.
proof. by rewrite big_cons big_nil addm0. qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_mkcond (P : 'a -> bool) (F : 'a -> t) s:
  big P F s = big predT (fun i => if P i then F i else idm) s.
proof.
  elim: s=> // x s ih; rewrite !big_cons -ih /predT /=.
  by case (P x)=> //; rewrite add0m.
qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_filter (P : 'a -> bool) F s:
  big predT F (filter P s) = big P F s.
proof. by elim: s => //= x s; case (P x)=> //; rewrite !big_cons=> -> ->. qed.

(* -------------------------------------------------------------------- *)
lemma nosmt eq_bigl (P1 P2 : 'a -> bool) (F : 'a -> t) s:
     (forall i, P1 i <=> P2 i)
  => big P1 F s = big P2 F s.
proof. by move=> h; rewrite /big (eq_filter h). qed.

(* -------------------------------------------------------------------- *)
lemma nosmt eq_bigr (P : 'a -> bool) (F1 F2 : 'a -> t) s:
     (forall i, P i => F1 i = F2 i)
  => big P F1 s = big P F2 s.
proof.                          (* FIXME: big_rec2 *)
  move=> eqF; elim: s=> // x s; rewrite !big_cons=> <-.
  by case (P x)=> // /eqF <-.
qed.

(* -------------------------------------------------------------------- *)
lemma big_andbC (P Q : 'a -> bool) F s:
  big (fun x => P x /\ Q x) F s = big (fun x => Q x /\ P x) F s.
proof. by apply/eq_bigl=> i. qed.

(* -------------------------------------------------------------------- *)
lemma nosmt eq_big (P1 P2 : 'a -> bool) (F1 F2 : 'a -> t) s:
     (forall i, P1 i <=> P2 i)
  => (forall i, P1 i => F1 i = F2 i)
  => big P1 F1 s = big P2 F2 s.
proof. by move=> /eq_bigl <- /eq_bigr <-. qed.

(* -------------------------------------------------------------------- *)
lemma nosmt congr_big r1 r2 P1 P2 (F1 F2 : 'a -> t):
     r1 = r2
  => (forall x, P1 x <=> P2 x)
  => (forall i, P1 i => F1 i = F2 i)
  => big P1 F1 r1 = big P2 F2 r2.
proof. by move=> <-; apply/eq_big. qed.

(* -------------------------------------------------------------------- *)
lemma big0_eq (P : 'a -> bool) s: big P (fun i => idm) s = idm.
proof.
  rewrite big_const; elim/Induction.natind: (count _ _) => n.
    by move/iter0<:t> => ->.
    by move/iterS<:t> => -> ->; rewrite addm0.
qed.

(* -------------------------------------------------------------------- *)
lemma big0 (P : 'a -> bool) F s: (forall i, P i => F i = idm) =>
  big P F s = idm.
proof. by move=> h; rewrite -@(big0_eq P s); apply/eq_bigr. qed.

(* -------------------------------------------------------------------- *)
lemma big_hasC (P : 'a -> bool) (F : 'a -> t) s: !has P s =>
  big P F s = idm.
proof.
  rewrite -big_filter has_count -size_filter.
  by rewrite ltz_def size_ge0 /= => /size_eq0 ->.
qed.

(* -------------------------------------------------------------------- *)
lemma big_pred0_eq (F : 'a -> t) s: big pred0 F s = idm.
proof. by rewrite big_hasC // has_pred0. qed.

(* -------------------------------------------------------------------- *)
lemma big_pred0 (P : 'a -> bool) (F : 'a -> t) s:
  (forall i, P i <=> false) => big P F s = idm.
proof. by move=> h; rewrite -@(big_pred0_eq F s); apply/eq_bigl. qed.

(* -------------------------------------------------------------------- *)
lemma big_cat (P : 'a -> bool) (F : 'a -> t) s1 s2:
  big P F (s1 ++ s2) = big P F s1 + big P F s2.
proof.
  rewrite !@(big_mkcond P); elim: s1 => /= [|i s1 ih].
    by rewrite @(big_nil P F) add0m.
  by rewrite !big_cons /(predT i) /= ih addmA.
qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_catl (P : 'a -> bool) (F : 'a -> t) s1 s2: !has P s2 =>
  big P F (s1 ++ s2) = big P F s1.
proof. by rewrite big_cat => /big_hasC ->; rewrite addm0. qed.
  
(* -------------------------------------------------------------------- *)
lemma nosmt big_catr (P : 'a -> bool) (F : 'a -> t) s1 s2: !has P s1 =>
  big P F (s1 ++ s2) = big P F s2.
proof. by rewrite big_cat => /big_hasC ->; rewrite add0m. qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_rcons (P : 'a -> bool) (F : 'a -> t) s x:
  big P F (rcons s x) = if P x then big P F s + F x else big P F s.
proof.
  by rewrite -cats1 big_cat big_cons big_nil; case: (P x); rewrite !addm0.
qed.

(* -------------------------------------------------------------------- *)
lemma nosmt eq_big_perm (P : 'a -> bool) (F : 'a -> t) s1 s2:
  perm_eq s1 s2 => big P F s1 = big P F s2.
proof.
  move=> /perm_eqP; rewrite !@(big_mkcond P).
  elim s1 s2 => [|i s1 ih1] s2 eq_s12.
    by case s2 eq_s12 => // i s2 h; cut := h (pred1 i); smt.
  cut r2i: mem s2 i by rewrite -has_pred1 has_count -eq_s12 smt.
  have/splitPr [s3 s4] ->> := r2i.
  rewrite big_cat !big_cons /(predT i) /=.
  rewrite addmCA; congr; rewrite -big_cat; apply/ih1=> a.
  by have := eq_s12 a; rewrite !count_cat /= addzCA => /addzI.
qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_rem (P : 'a -> bool) (F : 'a -> t) s x: mem s x =>
  big P F s = (if P x then F x else idm) + big P F (rem x s).
proof.
  by move/perm_to_rem/eq_big_perm=> ->; rewrite !@(big_mkcond P) big_cons.
qed.

(* -------------------------------------------------------------------- *)
lemma nosmt bigD1 (F : 'a -> t) s x: mem s x => uniq s =>
  big predT F s = F x + big (predC1 x) F s.
proof. by move=> /big_rem-> /rem_filter->; rewrite big_filter. qed.

(* -------------------------------------------------------------------- *)
lemma big_split (P : 'a -> bool) (F1 F2 : 'a -> t) s:
  big P (fun i => F1 i + F2 i) s = big P F1 s + big P F2 s.
proof.
  elim: s=> /= [|x s ih]; 1: by rewrite !big_nil addm0.
  rewrite !big_cons ih; case: (P x) => // _.
  by rewrite addmCA -!addmA addmCA.
qed.

(* -------------------------------------------------------------------- *)
lemma nosmt bigID (P : 'a -> bool) (F : 'a -> t) (a : 'a -> bool) s:
  big P F s = big (predI P a) F s + big (predI P (predC a)) F s.
proof.
  rewrite !@(big_mkcond _ F) -big_split /(/\) /[!].
  apply/eq_bigr => i _ /=.
  by rewrite /predI /predC; case: (a i); rewrite ?addm0 ?add0m.
qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_reindex (P : 'a -> bool) (F : 'a -> t) (f f' : 'a -> 'a) s:
  cancel f' f => big P F s = big (P \o f) (F \o f) (map f' s).
proof. by move=> bij_ff'; rewrite -big_map -map_comp id_map. qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_seq_cond (P : 'a -> bool) (F : 'a -> t) s:
  big P F s = big (fun i => mem s i /\ P i) F s.
proof. by rewrite -!@(big_filter _ _ s); congr; apply/eq_in_filter. qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_seq (F : 'a -> t) s:
  big predT F s = big (fun i => mem s i) F s.
proof. by rewrite big_seq_cond; apply/eq_bigl. qed.

(* -------------------------------------------------------------------- *)
lemma nosmt eq_big_seq (F1 F2 : 'a -> t) s:
     (forall x, mem s x => F1 x = F2 x)
  => big predT F1 s = big predT F2 s.
proof. by move=> eqF; rewrite !big_seq; apply/eq_bigr. qed.

(* -------------------------------------------------------------------- *)
lemma big1_eq (P : 'a -> bool) s: big P (fun x => idm) s = idm.
proof.
  rewrite big_const; elim/Induction.natind: (count _ _)=> n.
    by move/iter0<:t> => ->.
  by move/iterS<:t> => -> ->; rewrite addm0.
qed.

(* -------------------------------------------------------------------- *)
lemma big1 (P : 'a -> bool) F s:
  (forall i, P i => F i = idm) => big P F s = idm.
proof. by move/eq_bigr=> ->; apply/big1_eq. qed.

(* -------------------------------------------------------------------- *)
lemma big1_seq (P : 'a -> bool) F s:
  (forall i, P i /\ (mem s i) => F i = idm) => big P F s = idm.
proof. by move=> eqF1; rewrite big_seq_cond big_andbC big1. qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_int_cond m n P F:
  bigi P F m n = bigi (fun i => m <= i < n /\ P i) F m n.
proof. by rewrite /bigi big_seq_cond; apply/eq_bigl=> i /=; smt. qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_int m n F:
  bigi predT F m n = bigi (fun i => m <= i < n) F m n.
proof. by rewrite big_int_cond. qed.

(* -------------------------------------------------------------------- *)
lemma nosmt congr_big_int (m1 n1 m2 n2 : int) P1 P2 F1 F2:
       m1 = m2 => n1 = n2
    => (forall i, m1 <= i < n2 => P1 i = P2 i)
    => (forall i, P1 i /\ (m1 <= i < n2) => F1 i = F2 i)
    => bigi P1 F1 m1 n1 = bigi P2 F2 m2 n2.
proof.
  move=> <- <- eqP12 eqF12; rewrite /bigi big_seq_cond @(big_seq_cond P2).
  by apply/eq_big=> i /=; rewrite mem_range smt.
qed.

(* -------------------------------------------------------------------- *)
lemma nosmt eq_big_int (m n : int) F1 F2:
     (forall i, m <= i < n => F1 i = F2 i)
  => bigi predT F1 m n = bigi predT F2 m n.
proof. by move=> eqF; apply/congr_big_int. qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_ltn_cond (m n : int) P F: m < n =>
  let x = bigi P F (m+1) n in
  bigi P F m n = if P m then F m + x else x.
proof. by rewrite /bigi => /range_ltn ->; rewrite big_cons. qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_ltn (m n : int) F: m < n =>
  bigi predT F m n = F m + bigi predT F (m+1) n.
proof. by move/big_ltn_cond=> /= ->. qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_geq (m n : int) P F: n <= m =>
  bigi P F m n = idm.
proof. by move/range_geq; rewrite /bigi=> ->; rewrite big_nil. qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_addn (m n a : int) P F:
    bigi P F (m+a) n
  = bigi (fun i => P (i+a)) (fun i => F (i+a)) m (n-a).
proof.
rewrite /bigi range_addl big_map; apply/eq_big.
  by move=> i /=; rewrite /(\o) addzC.
by move=> i /= _; rewrite /(\o) addzC.
qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_int1 n F: bigi predT F n (n+1) = F n.
proof. by rewrite big_ltn 1:smt big_geq // addm0. qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_cat_int (n m p : int) P F: m <= n => n <= p =>
  bigi P F m p = (bigi P F m n) + (bigi P F n p).
proof. by move=> lemn lenp; rewrite {-1}/bigi -big_cat -range_cat. qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_int_recl (n m : int) F: m <= n =>
  bigi predT F m (n+1) = F m + bigi predT (fun i => F (i+1)) m n.
proof.
  move=> lemn; rewrite big_ltn 1?big_addn /= 1:smt.
  by rewrite subrE addrK.
qed.

(* -------------------------------------------------------------------- *)
lemma nosmt big_int_recr (n m : int) F: m <= n =>
  bigi predT F m (n+1) = bigi predT F m n + F n.
proof. by move=> lemn; rewrite @(big_cat_int n) ?big_int1 // smt. qed.