(* For the background behind these two functions (LSE and SLSE), see
  https://inria.hal.science/hal-0434315 *)

(* Preliminary definitions and lemmas *)
module Helpers
  use real.RealInfix
  use real.Abs
  use real.ExpLog
  use ufloat.USingle
  use ufloat.USingleLemmas
  use int.Int
  use real.Sum
  use real.FromInt

  function abs_fun (f:int -> real) : int -> real =
     fun i -> abs (f i)
  function abs_real_fun (f:int -> usingle) : int -> real =
    fun i -> abs_fun (real_fun f) i

  let rec lemma sum_of_fun_le_sum_of_abs_fun (f : int -> usingle) (a b:int)
    requires { 0 <= b-a }
    variant { b-a }
    ensures { abs (sum (real_fun f) a b) <=. sum (abs_real_fun f) a b }
  = if b-a > 0 then sum_of_fun_le_sum_of_abs_fun f a (b-1)

  let rec function usum_rec (f : int -> usingle) (a b:int) : usingle
    requires { 0 <= b - a }
    variant { b - a }
    ensures {
      abs (to_real result -. sum (real_fun f) a b)
       <=. sum (abs_real_fun f) a b *. (eps *. from_int (b-a))
    }
  = if b - a <= 0 then uzero else if b - a = 1 then f (b-1) else usum_rec f a (b-1) ++. f (b-1)

  constant log_error:real
  axiom log_error_bounds : 0. <=. log_error <=. 1.
  val function log_approx (x:usingle) : usingle
    requires { 0. <. to_real x }
    ensures {
      abs (to_real result -. log (to_real x)) <=. abs (log (to_real x)) *. log_error
    }

  constant log2_error:real
  axiom log2_error_bounds : 0. <=. log2_error <=. 1.
  val function log2_approx (x:usingle) : usingle
    requires { 0. <. to_real x }
    ensures {
      abs (to_real result -. log2 (to_real x)) <=. abs (log2 (to_real x)) *. log2_error
    }

  constant exp_error:real
  axiom exp_error_bounds : 0. <=. exp_error <=. 0x1p-3
  val function exp_approx (x:usingle) : usingle
    ensures {
      abs (to_real result -. exp (to_real x)) <=. exp_error *. exp (to_real x)
    }

  lemma exp_approx_pos : forall x. to_real (exp_approx x) >. 0.

  let rec ghost sum_strictly_pos (f: int -> real) (a b:int)
    variant { b - a }
    requires { a < b }
    requires { forall i. a <= i < b -> 0. <. f i }
    ensures { 0. <. sum f a b = sum (abs_fun f) a b }
  = if a < (b-1) then sum_strictly_pos f a (b - 1)

  let rec ghost usum_strictly_pos (f: int -> usingle) (a b:int)
    variant { b - a }
    requires { a < b }
    requires { forall i. a <= i < b -> 0. <. to_real (f i) }
    ensures { 0. <. to_real (usum_rec f a b) }
  = if a < (b-1) then usum_strictly_pos f a (b - 1)


  lemma lt_compat_order_mult_l : forall x y z. 0. <. z -> x <. y -> z *. x <. z *. y
end

module LSE
  use real.RealInfix
  use real.Abs
  use real.ExpLog
  use ufloat.USingle
  use ufloat.USingleLemmas
  use ref.Ref
  use int.Int
  use real.Sum
  use real.FromInt
  use Helpers

  use ufloat.HelperLemmas

  let ghost function exp_of_f (f: int -> usingle) = (fun i -> exp ((real_fun f) i))


  let lse (f : int -> usingle) (n:int) : usingle
    requires { 0 < n < 8388608 }
    ensures {
      let exact = log(sum (exp_of_f f) 0 n) in
      abs (to_real result -. exact) <=. abs(log(sum (exp_of_f f) 0 n)) *. log_error
      -. (log (1. -. (exp_error +. from_int n *. eps *. (1. +. exp_error)))) *. (1. +. log_error)
    }
  =
  let function f' = (fun i -> exp_approx (f i)) in
  let function exact_f' = (fun i -> (exp_of_f f) i) in
  assert {
   forall i. 0 <= i < n -> abs (to_real (f' i) -. (exact_f') i)
     <=. exp_error *. exact_f' i
  };
  let s = usum_rec f' 0 n in
  sum_single_error_propagation s f' exact_f' exact_f' (abs_real_fun f') n (eps *. from_int n) 0. exp_error 0.;
  assert {
    abs (to_real s -. sum (exact_f') 0 n)
    <=. (exp_error +. (from_int n *. eps) *. (1.0 +. exp_error)) *.
      (sum (exact_f') 0 n)
  };
  usum_strictly_pos f' 0 n;
  sum_strictly_pos (exp_of_f f) 0 n;
  lt_compat_order_mult_l (exp_error +. from_int n *. eps *. (1.0 +. exp_error)) 1. (sum exact_f' 0 n);
  log_approx s
end



module SLSE
  use real.RealInfix
  use real.Abs
  use real.ExpLog
  use ufloat.USingle
  use ufloat.USingleLemmas
  use int.Int
  use real.Sum
  use real.FromInt
  use Helpers

  constant max_size:int = 8388608

  constant a_eps:real = (eps /. 2.) *. (1. +. 4. *. (1. +. eps) *. (1. +. eps))
  constant a_max:real
  constant g_error:real = exp_error +. (exp (a_eps *. ((2. *. a_max +. 1.) *. (2. *. a_max +. 1.)) +. 1.5 *. eta) -. 1.)
     *. (1. +. exp_error)
  (* TODO: Instead bound a_max and prove a lemma on g_error_bound *)
  lemma g_error_bound : g_error <=. 0x1p-2

  let ghost function exact_g (x y mu:usingle) : real
  =
  exp (-.((to_real x +. to_real mu -. to_real y) *. (to_real x +. to_real mu -. to_real y)) /. 2.0)

  lemma compat_order_mult1 : forall x y z. 0. <=. z -> x <=. y -> x *. z <=. y *. z
  lemma compat_order_mult2 : forall x y z. 0. <=. z -> x <=. y -> z *. x <=. z *. y

  let function g (x y mu : usingle) =
    exp_approx ((--.((x ++. mu --. y) **. (x ++. mu --. y))) ///. utwo)

  lemma g_pos :
    forall x y mu. 0. <. to_real (g x y mu)

  let lemma error_of_g (x y mu : usingle)
    requires { abs (to_real x) <=. a_max }
    requires { abs (to_real y) <=. a_max }
    requires { 0. <=. to_real mu <=. 1. }
    ensures {
      abs (to_real (g x y mu) -. exact_g x y mu) <=. (exact_g x y mu) *. g_error
    }
  =
  let ret = exp_approx ((--.((x ++. mu --. y) **. (x ++. mu --. y))) ///. u2) in
  let t = abs (to_real x +. to_real mu) +. abs (to_real y) in
  let t' = 2. *. a_max +. 1. in
  let lemma err ()
  ensures {
    exact_g x y mu *. (exp_error +. (exp((eps +. (4. *. (eps +. eps *. eps)) *. (1. +. eps))
      *. ((t *. t) /. (to_real u2)) +. (eta /. to_real u2 +. eta)) -. 1.) *. (1. +. exp_error))
    <=. exact_g x y mu *. (exp_error +. (exp(a_eps *. (t' *. t') +. 1.5 *. eta) -. 1.) *. (1. +. exp_error))
  }
   =
  compat_order_mult1 t t' t;
  compat_order_mult2 t t' t';
  assert {
    exp((eps +. (4. *. (eps +. eps *. eps)) *. (1. +. eps))
      *. ((t *. t) /. (to_real u2)) +. (eta /. to_real u2 +. eta))
     <=. exp(a_eps *. (t' *. t') +. 1.5 *. eta)
  };
  compat_order_mult2
     (exp_error +. (exp((eps +. (4. *. (eps +. eps *. eps)) *. (1. +. eps))
       *. ((t *. t) /. (to_real u2)) +. (eta /. to_real u2 +. eta)) -. 1.) *. (1. +. exp_error))
     (exp_error +. (exp(a_eps *. (t' *. t') +. 1.5 *. eta) -. 1.) *. (1. +. exp_error))
     (exact_g x y mu);
  in
  assert {
  (-. (((to_real x +. to_real mu) -. to_real y)
          *. ((to_real x +. to_real mu) -. to_real y)))
     /. 2.
    <=. ((abs (to_real x +. to_real mu) +. abs (to_real y))
          *. (abs (to_real x +. to_real mu) +. abs (to_real y)))
         /. 2.
  };
  (* We call the strategy on the following assertion *)
  assert { abs (to_real ret -. exact_g x y mu) <=. (exact_g x y mu) *. g_error }


  let ghost function exact_f (a:int -> usingle) (mu : usingle) (n i : int) : real =
    log2(sum (fun j -> exact_g (a i) (a j) mu) 0 n)

  let function f (a:int -> usingle) (mu : usingle) (n i : int) : usingle
    requires { 0 <= i < n <= max_size }
    requires { forall k. 0 <= k < n -> abs (to_real (a k)) <=. a_max }
    requires { 0. <=. to_real mu <=. 1. }
    ensures {
      abs (to_real result -. exact_f a mu n i)
      <=. log2_error *. abs (exact_f a mu n i)
        -. log2 (1. -. (g_error +. from_int n *. eps *. (1. +. g_error))) *. (1. +. log2_error)
    }
  =
  let function g' = (fun j -> if 0 <= j < n then g (a i) (a j) mu else uone) in
  let ghost function exact_g' = (fun j -> exact_g (a i) (a j) mu) in
  let s = usum_rec g' 0 n in
  sum_single_error_propagation s g' exact_g' exact_g' (abs_real_fun g') n (eps *. from_int n) 0. g_error 0.;
  usum_strictly_pos g' 0 n;
  assert {
    exact_f a mu n i = log2 (sum exact_g' 0 n)
    by
     (fun (j:int) -> exact_g (a @ i) (a @ j) mu) = exact_g'
  };
  assert { sum exact_g' 0 n <> 0. };
  lt_compat_order_mult_l (g_error +. (eps *. from_int n) *. (1.0 +. g_error)) 1. (sum exact_g' 0 n);
  log2_approx (s)

  let function slse (a:int -> usingle) (mu : usingle) (n : int) : usingle
    requires { 0 < n <= max_size }
    requires { forall i. 0 <= i < n -> abs (to_real (a i)) <=. a_max }
    requires { 0. <=. to_real mu <=. 1. }
    ensures {
      let exact = sum (fun i -> exact_f a mu n i) 0 n in
      abs (to_real result -. exact)
      <=. (log2_error +. eps *. from_int n *. (1. +. log2_error))
            *. sum (fun i -> abs (exact_f a mu n i)) 0 n
          -. log2 (1. -. (g_error +. from_int n *. eps *. (1. +. g_error))) *. (1. +. log2_error)
           *. from_int n *. (1. +. eps *. from_int n)
    }
  =
  let function f' = (fun i -> if 0 <= i < n then f a mu n i else uzero) in
  let function exact_f' = (fun i -> exact_f a mu n i) in
  let function exact_f'' = (fun i -> abs (exact_f a mu n i)) in
  assert {
    forall i. 0 <= i < n ->
      abs (to_real (f' i) -. exact_f' i)
      <=. log2_error *. exact_f'' i
        -. log2 (1. -. (g_error +. from_int n *. eps *. (1. +. g_error)))
          *. (1. +. log2_error)
  };
  let ret = usum_rec f' 0 n in
  sum_single_error_propagation ret f' exact_f' exact_f'' (abs_real_fun f') n (eps *. from_int n) 0.
    log2_error (-. log2 (1. -. (g_error +. from_int n *. eps *. (1. +. g_error))) *. (1. +. log2_error));
  ret
end
