(* 
   O(train_bound * train_examples + test_bound * test_examples)
*)
let get_cumulatives (train_bound,test_bound) (train_examples, test_examples) (p : float) = 
  let train_cumulative = Array.create (train_examples + 1) Log_float.zero 
  and test_cumulative = Array.create (test_examples + 1) Log_float.zero in
  for i = 0 to train_examples do
    train_cumulative.(i) <- train_bound i p
  done;
  for i = 0 to test_examples do
    test_cumulative.(i) <- test_bound i p
  done;
  train_cumulative, test_cumulative

exception Done of int

(* for increasing cumulative distributions 
   O(training_examples)
*)
let find_thresh_down (threshold : float) cumulative start =
  try
    for i = start downto 0 do 
      if cumulative.(i) <= threshold then raise (Done i)
    done; ~- 1;
  with Done i -> i

(* explicitly calculate the (bound on the) expected value of the size of the confidence interval 
   O(training_examples + test_examples)
*)
let get_sum_up (train_cumulative, test_cumulative) (train_dist,test_dist) threshold =
  let sum = ref Log_float.zero in
  let thresh_point = ref (Array.length train_cumulative - 1) in
  try
    for j = 0 to Array.length test_cumulative - 1 do
      let thresh = find_thresh_down (threshold -. test_cumulative.(j)) train_cumulative !thresh_point in
      if thresh >= 0 then 
	begin 
	  sum := Log_float.add (train_cumulative.(thresh) +. test_dist.(j)) !sum;
	  thresh_point := thresh;
	end
      else raise (Done 0)
    done; !sum
  with Done i -> !sum
      
(* for decreasing cumulative distributions 
   O(training_examples)
*)
let find_thresh_up (threshold : float) cumulative start =
  try
    for i = start to Array.length cumulative - 1 do 
      if cumulative.(i) <= threshold then raise (Done i)
    done; Array.length cumulative;
  with Done i -> i

(* 
   Explicitly calculate the (bound on the) expected value of the size of the confidence interval 
   O(training_examples + test_examples)
*)
let get_sum_down (train_cumulative, test_cumulative) (train_dist,test_dist) threshold =
  let sum = ref Log_float.zero in
  let thresh_point = ref 0 in
  try 
    for j = Array.length test_cumulative - 1 downto 0 do
      let thresh = find_thresh_up (threshold -. test_cumulative.(j)) train_cumulative !thresh_point in
      if thresh < Array.length train_cumulative 
      then begin 
	sum := Log_float.add (train_cumulative.(thresh) +. test_dist.(j)) !sum;
	thresh_point := thresh;
      end
      else raise (Done 0)
    done;
    !sum
  with Done i -> !sum

(*
   O(train_bound * train_examples + test_bound * test_examples)
*)
let bound_upper bounds examples (train_error,test_error) (p : float) = 
  let train_cumulative, test_cumulative = get_cumulatives bounds examples p in
  let threshold = test_cumulative.(test_error) +. train_cumulative.(train_error) in

  let train_dist = Array.copy train_cumulative
  and test_dist = Array.copy test_cumulative in
  for i = 1 to Array.length train_cumulative - 1 do
    train_dist.(i) <- Log_float.sub train_cumulative.(i) train_cumulative.(i-1)
  done;
  for i = 1 to Array.length test_cumulative - 1 do
    test_dist.(i) <- Log_float.sub test_cumulative.(i) test_cumulative.(i-1) 
  done;
  
  get_sum_up (train_cumulative, test_cumulative) (train_dist, test_dist) threshold 

(*
   
*)
let bound_lower bounds examples (train_error,test_error) (p : float) = 
  let train_cumulative, test_cumulative = get_cumulatives bounds examples p in
  let threshold = test_cumulative.(test_error) +. train_cumulative.(train_error) in

  let train_dist = Array.copy train_cumulative
  and test_dist = Array.copy test_cumulative in
  for i = 0 to Array.length train_cumulative - 2 do
    train_dist.(i) <- Log_float.sub train_cumulative.(i) train_cumulative.(i+1);
  done;
  for i = 0 to Array.length test_cumulative - 2 do
    test_dist.(i) <- Log_float.sub test_cumulative.(i) test_cumulative.(i+1);
  done;
  
  get_sum_down (train_cumulative, test_cumulative) (train_dist, test_dist) threshold 
