Solution for coursera's Cryptography programming task 1.1

task_1_1.ml

open Batteries_uni

let regroup e n =
  let next () =
    if Enum.is_empty e then
      raise Enum.No_more_elements
    else
      Enum.take n e in Enum.from next;;

let load_ciphers fname =
  let lines = File.lines_of fname in
  let process_line line =
    let regroupped = regroup (String.enum line) 2 in
    let get_ord x = int_of_string ("0x" ^ (String.of_enum x)) in
    Enum.map get_ord regroupped in
  Enum.map process_line lines;;

(* Should be fine for small files *)
let load_file fname =
  File.with_file_in fname (fun f -> String.enum (BatInnerIO.read_all f));;

let get_histogram fname =
  let arr = Array.make 256 0 in
  let count = ref 0 in
  let process_char c = let x = int_of_char c in begin
    arr.(x) <-arr.(x) + 1;
    count :=  !count + 1;
  end in
  (Enum.iter process_char (load_file fname);
  (arr,!count));;

let (hist, samples_count) = get_histogram "shakespeare.txt";;

let (msg, ciphers) =
  let e = load_ciphers "ciphers.txt" in
  let h = Array.of_enum (Option.get (Enum.get e)) in
  let t = List.of_enum (Enum.map Array.of_enum e) in
  (h,t);;

let get_p x =
  let p = hist.(x) in
  (float_of_int p) /. (float_of_int samples_count);;

let calculate_probs msg ciphers =
  let msg_len = Array.length msg in
  let get_val x arr = try Array.get arr x with Invalid_argument _ -> 1 in
  (* i is msg index(0--msg_len-1), j is 0--255 *)
  let pseudo_p i j =
    let multiply_p p cipher =
      let m = ((get_val i msg) lxor (get_val i cipher)) in
      p *. (get_p (j lxor m)) in
    (get_p j) *. List.fold_left multiply_p 1. ciphers in
  let unnormalized = Array.of_enum (Enum.map (fun i -> Array.of_enum (Enum.map (pseudo_p i) (0--255))) (0--(msg_len-1))) in
  unnormalized;;

let get_maxprop_string probs =
  let find_max arr =
    let m = Array.max arr in
    Array.findi (fun x -> 0 == compare x m) arr |> char_of_int in
  Array.map find_max probs |> Array.enum |> String.of_enum;;

let main () =
  let s = calculate_probs msg ciphers |> get_maxprop_string in
  print_endline s;;

let () = main ();;

(* P'[m0=i|m0+m1=b,m0+m2=c]=P[m0+m1=b,m0+m2=c|m0=i]*P[m0=i]=P[m1=b+i]*P[m2=c+i]*P[m0=i]*)