(* -------------------------------------------------------------------- *)
require import AllCore.
require import Distr List Aprhl StdRing StdOrder StdBigop.
(*---*) import IntOrder RField RealOrder.
(* -------------------------------------------------------------------- *)
op eps: { real | 0%r <= eps } as ge0_eps.

hint exact : ge0_eps.


op N : { int | 0 < N } as gt0_N.

hint exact : gt0_N.


pred adj: int list & int list.

op evalQ : int -> int list -> int.

axiom one_sens d1 d2 i: adj d1 d2 => `|evalQ  i d1 - evalQ i d2| <= 1.

(* -------------------------------------------------------------------- *)

(* helper lemma ------------------------------------------------------- *)

lemma sumr0 (n : int) :
  0 <= n => sumr n (fun (_ : int) => 0%r) = 0%r.
proof.
move => ge0_n.
rewrite /sumr sumri_const 1:ge0_n 1:intmulr /#.
qed.

lemma count_iota_0 : forall n,
count (transpose (=) n) (iota_ 0 n) = 0.
    proof.
    move => n.
      apply count_eq0. rewrite hasPn. 
    move => x.
    smt(mem_iota).
qed.

lemma count_iota : forall n j,
n=size(iota_ 0 n) =>  0<=j< n =>  (count (fun (i : int) => j = n - i - 1) (iota_ 0 n)) = 1.
    proof.
    elim/natind.
    move => n Hn j Hs Hj. smt(size_iota).
    move => n Hn HI j Hs Hj. 
      rewrite iotaSr. assumption. simplify.  rewrite -cats1. rewrite count_cat. 
      rewrite /count. case (j=0) => H. rewrite H.
      simplify. 
     have H0 : (n+1-n-1=0). smt(). 
      rewrite H0. simplify.
      have H1 :(fun i => i = n) =(fun i => 0=n+1-i-1).  smt(). 
      rewrite -H1. rewrite count_iota_0. smt(). 
      have H0 : (n+1-n-1=0). smt().
      rewrite H0 H. simplify.    
      have H1 :(fun i => j-1 = n-i-1) =(fun i => j=n+1-i-1).  smt().
      rewrite -H1. rewrite HI.   smt(size_iota). smt(). smt(). 
   qed.

  lemma count_iota0 : forall n j,
!0<=j< n =>  (count (fun (i : int) => j = n - i - 1) (iota_ 0 n)) = 0.
    proof.
    elim/natind.
    move => n Hn j Hs.
    search count.
      rewrite -count_eq0.
      smt.
    move => n Hn HI j  Hj.
   rewrite iotaSr. assumption. simplify.  rewrite -cats1. rewrite count_cat. 
      rewrite /count. case (j=0) => H. rewrite H.
      simplify. 
     have H0 : (n+1-n-1=0). smt(). 
      rewrite H0. simplify.
      have H1 :(fun i => i = n) =(fun i => 0=n+1-i-1).  smt(). 
      rewrite -H1. rewrite count_iota_0. smt(). 
      have H0 : (n+1-n-1=0). smt().
      rewrite H0 H. simplify.    
      have H1 :(fun i => j-1 = n-i-1) =(fun i => j=n+1-i-1).  smt().
      rewrite -H1. rewrite HI.   smt(size_iota). smt(). 
  qed.


lemma bigi_eps : forall  n j,
    0<= j < n=>
    eps = bigi (fun (i : int) => j = n-i -1 ) (fun (_ : int) => eps) 0 n.
    proof.
    move => ? ? ?. rewrite big_const. rewrite count_iota.   smt(size_iota).
      assumption. smt(iter1). 
  qed.
  

  lemma bigi_eps0 : forall  n j,
  ! ( 0<= j < n)=>
    0%r = bigi (fun (i : int) => j = n-i -1 ) (fun (_ : int) => eps) 0 n.
    proof.
    move => ? ? ?.  rewrite big_const. rewrite count_iota0.
      trivial.  smt(iter0).
  qed.


lemma bigop : forall x,
   (bigi predT (fun (i : int) => if x = N - i - 1 then eps else 0%r) 0 N)<= eps .
    proof.
    move => x.
    rewrite -big_mkcond.
    case (0 <= x < N). 
    + move => ?.
      have H0 : (eps = bigi (fun (i : int) => x = N - i - 1) (fun (_ : int) => eps) 0 N).
      apply (bigi_eps) => //.
      rewrite -H0 => //.
    + move => ?.
      rewrite -bigi_eps0 => //.
qed.  

(* -------------------------------------------------------------------- *)

(********Report Noise Max Algorithms********)
(* GOAL *)
module M = {
  var db : int list 
  proc rnm() : int = {
    var cur  : int;
    var i  : int;
    var output : int;
    var max : int;
    i  <- 0;
    max <- 0;
    output <- -1;
    cur <- 0;
    while (i < N) {
        cur <$ lap eps (evalQ i db);
        if (max < cur) {max <- cur; output <- i;}
        i <- i + 1;   
    }
    return output;
  }
}.
  
(* Main Lemma ------------------------------------------------------- *)

lemma rnm : aequiv [ [ eps & 0%r ]
  M.rnm ~ M.rnm :(adj M.db{1} M.db{2}) ==> ={res}].
proof.
  proc.
  seq 4 4 :
    (adj M.db{1} M.db{2} /\ ={i, output, max, cur} /\
     i{1}=0 /\ max{1}=0 /\ cur{1}=0). 
  toequiv; auto.
  pweq (output, output).
    while true (N - i).  (* N - i is for showing termination *)
    move => z. wp. rnd predT. skip. progress.
    by rewrite lap_ll.
    smt().  (* termination metric goes down *)
    skip; progress. 
    (* when termination metric is <= 0 => boolean expression is false *)
    smt().
    while true (N - i); auto; smt(lap_ll).  (* auto figures out the right pred *)
    smt().
    move => x.
    (* use Consequence Rule to say when we spend our budget *)
    conseq
      <[ (bigi predT (fun (i : int) => if x = N - i - 1 then eps else 0%r) 0 N) &
         0 %r]>.
    move => _ _ /> _ . apply bigop.
    awhile [ (fun i  => if (x = N - i - 1) then eps else 0%r ) &
             (fun _ => 0%r) ]
           N [N - i - 1] 
      (adj M.db{1} M.db{2} /\ ={i, output}
         /\ (max{1} < cur{1} => output{1} = i{1})
         /\ (max{2} < cur{2} => output{1} = i{1})).
      rewrite ltzW gt0_N.
      smt(ge0_eps).
      smt().
      smt().
      by rewrite sumr0 1:ltzW.
      move => k.
        have cases : (x =  N - k -1) \/ (x <>  N - k -1).
          by case (x = N - k - 1).
        (* split into two cases: *)
        elim cases => [eq_x_N_min_k_min1 | ne_x_N_min_k_min1].
        seq 1 1 :  (* handle lap *)
          (adj M.db{1} M.db{2} /\ ={i, output} /\
           (max{1} < cur{1} => output{1} = i{1}) /\
           (max{2} < cur{2} => output{1} = i{1}) /\
           i{1} < N /\ i{2} < N /\ k = N - i{1} - 1 /\ N - i{1} - 1 <= N /\
           ={cur})  (* NOTE - because first arg to lap is 0 *)
          <[ (if x = N - k - 1 then eps else 0%r) & 0%r ]>.
        lap 0 1. smt(). smt(one_sens).
        if{1}.
        if{2}.
        toequiv; auto; smt().
        toequiv; auto; progress.
        by rewrite H0.
        trivial.  (* contradiction in assumptions *)
        smt().
        if{2}.
        toequiv; auto; progress.
        by rewrite H1.
        trivial.  (* contradiction in assumptions *)
        smt().
        toequiv; auto; smt().
        (* second case *)
        (* bring evalQ i M.db in {1} and {2} into assumptions: *)
        exists* (evalQ i M.db){1}; elim* => e1;
        exists* (evalQ i M.db){2}; elim* => e2.
        seq 1 1 : (* handle lap *)
          (e1 = evalQ i{1} M.db{1} /\ e2 = evalQ i{2} M.db{2} /\
           adj M.db{1} M.db{2} /\ ={i, output} /\
           (max{1} < cur{1} => output{1} = i{1}) /\
           (max{2} < cur{2} => output{1} = i{1}) /\
           i{1} < N /\ i{2} < N /\ x <> N - k - 1 /\
           k = N - i{1} - 1 /\ N - i{1} - 1 <= N /\
           (* NOTE - because first arg to lap is (e2 - e1): *)
           cur{1} + (e2 - e1) = cur{2})  
          <[ (if x = N - k - 1 then eps else 0%r) & 0%r ]>.
        lap (e2 - e1) 0.
        smt().
        progress. smt().
        if{1}.
        if{2}.
        toequiv; auto; progress. smt().
        toequiv; auto; progress.
        by rewrite H0.
        trivial.  (* contradiction in assumptions *)
        smt().
        if{2}.
        toequiv; auto; progress.
        by rewrite H1.
        trivial.  (* contradiction in assumptions *)
        smt().
        toequiv; auto; progress.
        trivial.  (* contradiction in assumptions *)
        trivial.  (* contradiction in assumptions *)
        smt().        
qed.