(* entropy.ml - Don Yang (uguu.org)

   02/17/03
                                                                           *)

(* Open first argument as input stream if given, otherwise use stdin *)
let input () =
   if Array.length Sys.argv > 1
   then open_in_bin Sys.argv.(1)
   else stdin;;

(* Get frequencies -- recursive version, overflows stack *)
let rec countfreq instream f =
   try
      let x = input_byte instream in
      f.(x) <- f.(x) + 1; 1 + (countfreq instream f)
   with End_of_file -> close_in instream; 0;;

(* Get frequencies -- iterative version (lame), doesn't overflow stack *)
let countfreq instream f =
   let x = [|0; 0|] in
   while
      (
         x.(0) <-
         (
            try input_byte instream
            with End_of_file -> close_in instream; -1
         );
         x.(0)
      ) != -1
   do
      f.(x.(0)) <- f.(x.(0)) + 1;
      x.(1) <- x.(1) + 1;
   done;
   x.(1);;

(* Compute entropy *)
let rec entropy f s i =
   if i > 255
   then 0.0
   else
      entropy f s (i + 1) -.
      (
         if f.(i) > 0
         then ((float f.(i)) /. s) *. log((float f.(i)) /. s) /. log(2.)
         else 0.0
      );;


(* Program entry point *)
let freq = Array.make 256 0;;
let filesize = (countfreq (input ()) freq);;
print_float(entropy freq (float filesize) 0);;
print_newline();;